简介：这是一套面向毕业设计/课程设计的机器学习危险物品识别项目，聚焦充电宝检测场景。项目完整交付代码与数据集，训练集覆盖带电芯充电宝与不带电芯充电宝两类样本，按1:10比例分布（500:5000），并将类别不均衡下模型偏好问题作为训练重点，提供训练集/测试集划分与mAP评估流程。资料包共50个文件，以Python脚本（26个py）为骨架，覆盖数据预处理、图像处理、模型训练、测试评估等环节；同时包含配置txt、预训练权重pkl、shell运行脚本、可视化示例图、训练记录以及报告文档，压缩包整体14.61MB，目录模块划分清晰。已有264人学习下载。项目依托SSD目标检测框架，提供完整工程目录、演示demo与说明文档，附带多个epoch的评估输出，便于读者复现训练流程、理解样本不均衡问题的处理思路，并迁移到其他危险品识别任务中。
1. 充电宝危险品识别的完整工程：值得先跑通再谈优化
做安检图像识别的人应该都有同感：真正把项目拖垮的往往不是模型结构，而是数据。这份资源打包了一个完整的充电宝危险品识别工程，从 SIXray 数据集的解析、VOC 格式转换，到 SSD300 的模型定义、训练、测试和 mAP 评估，全部跑通。它最有价值的地方在于，训练数据里带电芯充电宝与不带电芯充电宝的比例是 1:10，总共 5500 张图里正样本只有 500 张——这正是大多数毕设和课设最容易翻车的样本不均衡场景。适合拿来做毕业设计、课程设计参考，也适合想快速看一遍危险品检测完整流程的学习者。
2. 先把数据捋顺：从 SIXray 原始标注到可训练的 VOC 格式
2.1 认识这份代码里的数据管线
解压之后，你会看到很多文件名，第一次接触时容易懵。和我之前拆过的很多检测工程一样，这些文件的职责可以归成三条线：数据解析、图像处理、训练评估。
先看数据解析这条线。SIXray.py负责读取 SIXray 数据集的标注信息，这个数据集本身是安检场景的 X 光图像，标注格式和常规的 PASCAL VOC 不完全一样。voc0712.py和coco.py是 SSD 训练时常用的数据加载入口，前者读 VOC 格式，后者读 COCO 格式。这里的工程实际是把 SIXray 的数据转成了 VOC 风格，再用voc0712.py喂给训练脚本。
再看图像处理这条线。data_handler.py和image_handler.py负责图片的读取、尺寸调整、边界框坐标的换算。handler_5000.py只看名字就知道是处理那 5000 张不带电芯充电宝样本的脚本。这种按数量拆分的做法很实用，因为正负样本比例悬殊，分开处理能避免在代码里反复判断类别。
最后是模型和训练这条线。ssd.py是模型定义，ssd300_120000应该是预训练或已训练好的模型目录，train.py负责训练，test.py和eval.py负责推理和评估，eval_model145epoch_500和eval_model150epoch_500是两个训练后的权重目录，epoch 数分别是 145 和 150，后面那个 500 可能是每轮参与训练的样本数或验证集大小。
2.2 数据划分脚本的实际逻辑
train_test_txt.py这个脚本的核心作用，是把图像路径和标注信息写进训练集和测试集的 txt 文件。这种 txt 驱动的方式在很多检测框架里都常见，比直接在代码里写死路径更灵活。
下面是一段按照它的逻辑简化的划分脚本，用来理解它的工作方式：
import os import random from sklearn.model_selection import train_test_split image_dir = './data/images' annotation_dir = './data/annotations' all_images = [f for f in os.listdir(image_dir) if f.endswith('.jpg')] # 生成标注文件路径，VOC 格式下同名 xml 放在 annotation 目录 all_annotations = [os.path.join(annotation_dir, img.replace('.jpg', '.xml')) for img in all_images] # 按 8:2 划分训练集和测试集， stratify 参数保证正负样本比例在划分前后一致 train_imgs, test_imgs, train_anns, test_anns = train_test_split( all_images, all_annotations, test_size=0.2, stratify=[get_label(img) for img in all_images], # 用类别标签做分层 random_state=42 ) with open('./train.txt', 'w') as f: for img, ann in zip(train_imgs, train_anns): f.write(f'{os.path.join(image_dir, img)} {ann}\n') with open('./test.txt', 'w') as f: for img, ann in zip(test_imgs, test_anns): f.write(f'{os.path.join(image_dir, img)} {ann}\n')这里的stratify参数是关键。因为带电芯和不带电芯样本是 1:10，如果不做分层抽样，随机划分很容易出现测试集里正样本只剩几十张的情况，mAP 算出来方差极大。加了分层之后，训练集和测试集都保持约 1:10 的比例，训练集里正样本约 400 张，测试集约 100 张，评估结果更稳定。
random_state=42固定随机种子，保证每次跑出来的划分结果一致。这对复现实验特别重要，不然每次训练前数据划分一变，数据分布跟着变，你根本分不清模型效果变化是来自参数调整还是来自数据变化。
2.3 数据增强：对抗类别偏好的第一道防线
augmentations.py里定义了训练时的数据增强策略，这份增强不是随便加几个亮度变换就算完。检测任务里常用的 PhotometricDistort（光度失真）、RandomSampleCrop（随机裁剪）、RandomMirror（随机翻转）在这里都能看到作用。
我一般会在augmentations.py里重点看两个地方：一是是否对图像做了归一化，二是随机裁剪时如何处理超出了图像边界的框。很多检测代码在裁剪后没有把超出边界的坐标裁回图像范围内，导致训练时 loss 直接爆炸，或者出现接近零的 ground truth 框。
对于这种 1:10 的不均衡数据，增强策略应该对占比 10% 的类别有所倾斜。常见做法是，在每次迭代时以 50% 的概率对正样本做额外的随机旋转和亮度扰动，对负样本保持常规增强。这份项目代码里虽然没有单独写正样本过采样逻辑，但handler_5000.py这类脚本的存在，说明原始数据的预处理阶段已经做了大量负样本筛选。你可以在此基础上，把负样本从 5000 张里随机抽出一部分参与每轮训练，相当于在训练过程中动态调整正负样本比例。
数据增强的实际效果是加宽了正样本的外观变化范围。充电宝在 X 光图像里的形状差异其实不小，有的带外接充电线，有的外壳纹理不同，单靠 500 张原图很难覆盖。通过随机裁剪、旋转、颜色扰动，相当于把 500 张扩成几千个有效样本，这是解决类别偏好问题的第一道防线。
3. 模型训练与评估：SSD300 的完整闭环
3.1 config.py 中的关键参数
config.py是 SSD 训练里的全局配置，打开这个文件你会看到一堆数字，但它们并不是随便填的。之前我拆过不少 SSD 工程，这里面的参数直接影响训练效果。
# config.py 关键参数解析 EXCLUDE = ['__background__', 'person', 'bicycle', 'car'] # 不需要的类别 TRAIN_DATA = 'data/voc0712_train.txt' # 训练集路径索引 VAL_DATA = 'data/voc0712_test.txt' # 测试集路径索引 NUM_CLASSES = 4 # 背景 + 带电芯充电宝 + 不带电芯充电宝 + 其他危险品 BATCH_SIZE = 8 NUM_EPOCHS = 150 LEARNING_RATE = 1e-3 MILESTONES = [80, 120] # 在第 80、120 个 epoch 时降低学习率NUM_CLASSES这里不要只写 2，因为 SSD 的类别数要加上背景类，实际上输出通道数等于类别数加 4 倍的先验框数。如果这里写错，模型定义阶段的num_classes和数据集返回的标签维度对不上，训练时就会报维度的张量形状错误。
MILESTONES是学习率衰减的节点。前 80 个 epoch 用 1e-3 让模型快速收敛，80 到 120 降到 1e-4 精调，120 之后用 1e-5 微调。这个设置在 150 个 epoch 的训练中比较合理。如果你发现训练 loss 下降很快但 mAP 上不去，可以先检查是不是学习率降得太晚或太早。
3.2 跑通 train.py：从启动到保存模型
train.py的启动命令一般长这样，训练日志和模型权重都会输出到指定目录：
python train.py \ --config config.py \ --dataset_type voc0712 \ --data_root ./data \ --batch_size 8 \ --num_epochs 150 \ --save_dir ./weights如果你在自己机器上跑，可能会发现--batch_size设到 16 就会爆显存，这是正常的，因为 SSD300 的输入尺寸是 300x300，但每张图会生成 8732 个先验框，中间特征层的显存占用不小。8 是比较稳妥的起始值。显存足够的话可以试着调到 16，收敛速度会有提升，但 mAP 的最终值不一定更高。
训练过程中，我习惯每 5 个 epoch 在验证集上跑一次 mAP，而不是等全部训练结束才看结果。这样做的目的在于，尽早发现模型是否过拟合到少数类上。如果验证集 mAP 在 100 个 epoch 后开始下降，但训练集 loss 还在往下走，这就是典型的过拟合信号。这份工程里的eval_model145epoch_500和eval_model150epoch_500两个权重目录，说明原作者也是每隔一段时间做一次评估，把表现最好的权重保存下来。
train.py里通常会有保存逻辑：
# 训练循环中定期保存模型 if epoch % 5 == 0: torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': avg_loss, }, f'./weights/ssd300_epoch_{epoch}.pth')这里保存的不只是模型权重，还有优化器状态。如果你中途中断训练想恢复，直接加载这个文件就能接着跑。我踩过一次坑：只保存model_state_dict，跳过优化器状态，恢复训练后学习率被重置，效果直接崩掉。
3.3 eval.py 与 mAP 的严谨计算
eval.py里的 mAP 计算逻辑，决定了最终报告的数值能不能让人信服。项目摘要里明确要求划分测试集并计算测试集上的 mAP，所以这里要看得细一些。
# eval.py 中计算 mAP 的核心流程 def evaluate(model, test_loader, iou_threshold=0.5): all_detections = [] # 保存所有图片的检测结果 all_ground_truths = [] for images, targets in test_loader: with torch.no_grad(): outputs = model(images) # outputs 包含 predicted boxes, scores, labels batch_boxes, batch_scores, batch_labels = outputs # 按图片分别处理，将检测结果存入列表 for i in range(images.size(0)): keep = batch_scores[i] > 0.01 # 低分框先过滤 boxes = batch_boxes[i][keep] scores = batch_scores[i][keep] labels = batch_labels[i][keep] # 对每个类别分别计算 AP for cls in range(1, NUM_CLASSES): cls_mask = labels == cls cls_boxes = boxes[cls_mask] cls_scores = scores[cls_mask] # 计算 precision-recall 曲线下的面积 ap = compute_ap(cls_boxes, cls_scores, targets[i]['boxes'][targets[i]['labels'] == cls], iou_threshold) print(f'Class {cls} AP: {ap:.4f}')mAP 是所有类别 AP 的平均值。这里的iou_threshold=0.5是标准设定，但要注意：如果你的测试集里目标普遍很小，0.5 的 IoU 会显得过于严格，导致 mAP 偏低。可以对小目标放宽到 0.3 单独看一个版本。
低分框过滤的0.01也很关键。SSD 会输出海量低置信度框，如果不提前过滤，后面的 NMS 会非常慢。这个阈值设得太高会把一些置信度不高但实际的检测结果丢掉，设得太低又会让计算量陡增。0.01 是 SSD 代码里的经典默认值。
4. 避坑/常见问题：样本不均衡和复现中的血泪经验
4.1 模型输出清一色的不带电芯充电宝
现象：训练完 150 个 epoch，测试时发现模型把几乎所有观测都判为不带电芯充电宝，带电芯的识别率几乎为零。
原因：这是典型的类别偏好。两个类别比例 1:10，模型在训练过程中发现只需要把一切预测为多数类就能得到很低的 loss，因为这个方向的梯度惩罚在整体中占比过低。
解决：先做分层划分保证测试集比例不失真，再在训练中加入正样本过采样。具体做法是每次迭代时，从 500 张带电芯样本里随机抽取 100 张，和 500 张不带电芯样本混合成一个 mini-batch，这样正负比变成 1:5 而不是 1:10。损失函数中的正负样本权重也要调整，给少数类更高的 loss 权重。
4.2 训练集 mAP 很高、测试集 mAP 很低
现象：同一个权重在训练集上 mAP 能到 85% 以上，拿到测试集上直接掉到 50% 以下。
原因：训练过程中，数据增强做了随机裁剪和颜色抖动，但测试脚本里没有做同样的处理，而且测试时图片的缩放方式和训练不一致。比如训练时用的是RandomSampleCrop，评价时直接Resize到 300x300，目标在图片中的尺度和位置分布完全不同。
解决：检查test.py或eval.py中的预处理部分，确保和train.py的预处理一致，尤其是归一化的均值和标准差、颜色通道顺序、图片缩放方式。我在自己项目里曾经因为训练用 BGR、测试用 RGB 输出，导致 mAP 从 78% 掉到 30%，排查了一整个下午才发现是通道顺序问题。
4.3 加载预训练权重时 key 不匹配
现象：使用ssd300_120000预训练权重时，报出Missing key(s) in state_dict或者尺寸不匹配的错误。
原因：预训练权重的类别数和当前任务的类别数不一致。SSD 输出的分类分支最后的卷积层通道数是(类别数 + 1) * 24，类别数变了，这个卷积层的权重形状就不同。
解决：加载权重时，只加载 backbone 部分的权重，不加载最后的分类和回归分支。常见的做法是：
pretrained_dict = torch.load('ssd300_120000.pth') model_dict = model.state_dict() pretrained_weights = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].shape == v.shape} model_dict.update(pretrained_weights) model.load_state_dict(model_dict)这样处理之后，分类分支从头训练，backbone 部分继承预训练特征。如果你直接整体加载，训练一开始 loss 就可能是 NaN。
4.4 训练时内存或显存占用不断上涨导致死机
现象：训练到第 40 个 epoch 时，显存占用从 8GB 涨到 12GB，随后程序卡死。
原因：数据加载器DataLoader里的num_workers设得过大，加上训练循环里没有清理中间变量，GPU 显存被逐步占满。
解决：先把num_workers设为 2 或 4，再确认训练循环里的outputs、loss等张量在每轮迭代结束后正确释放。如果用的是 PyTorch，可以在每轮结束加torch.cuda.empty_cache()，但更根本的办法是把 batch_size 降低到 4，并检查requires_grad设置。还有一个常见坑：把验证集的数据也放到 GPU 上参与训练，但验证数据本身不做梯度更新，这也会让显存多占一份。
4.5 训练的损失值一直在震荡不下降
现象：loss 在前 10 个 epoch 里从 8 降到 5，之后就一直在 4.5 到 5.5 之间来回震荡，不再下降。
原因：学习率过大，或者数据增强太激进，导致模型在最优解附近来回跳动。有时又反过来是学习率太小，模型更新幅度不够。
解决：把学习率从 1e-3 降到 3e-4 再看。如果降了之后 loss 开始稳步下降，说明原来就是学习率问题。如果还是不降，那就把数据增强里的随机裁剪概率调低一些。
5. 用训练好的权重快速验证新图片，并迁移到自己的数据集
拿到eval_model150epoch_500这个权重后，怎么在单张图片上做推理，往往比重新训练一遍更常用。demo目录里的live.py或demo.ipynb能帮上忙。推理脚本的套路是这样的：
from ssd import build_ssd import torchvision.transforms as transforms from PIL import Image # 构建模型并加载权重 model = build_ssd('test', 300, 4) # 4 对应背景 + 3 个类别 model.load_state_dict(torch.load('eval_model150epoch_500/ssd300_150.pth')) model.eval() # 读取图片并预处理 image = Image.open('example.jpg').convert('RGB') transform = transforms.Compose([ transforms.Resize((300, 300)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) input_tensor = transform(image).unsqueeze(0) with torch.no_grad(): detections = model(input_tensor) # detections 的维度是 (1, num_classes, top_k, 5) # 最后一个维度的 5 个数分别是 score, cx, cy, w, h在跑这段代码前，务必确认权重里的类别顺序和你定义的一致。很多人在这一步翻过车：用自己训练的权重但忘记修改类别列表，导致第一个类输出的是背景概率，画出来的框全错位。
把这份工程迁移到自己的数据集上时，要改的文件很少：config.py里的类别数和类别名，voc0712.py里的数据读取路径，以及最终的输出类别映射。如果你要识别的目标远小于 300x300 里的常规尺度，可以把输入图片 resize 成 512x512，但对应的 SSD 模型也要改成 SSD512 结构，不能直接套用 300 的模型代码。
我自己后来做危险品识别时，每次换数据集都强制走一遍这几个步骤：先用分层划分脚本生成正则的 txt，再在训练初期用 20 个 epoch 快速验证 loss 是否稳定，确认之后再跑完整训练，最后用 0.5 的 IoU 阈值和 0.01 的置信度阈值做标准评估。这一套流程让我少踩了很多坑，希望帮到你。
本文还有配套的精品资源，点击获取
