基于FasterRCNN目标检测的缺陷检测算法 数据集包含五类别具体如图所示 共计1800张图 包含VOC格式数据集Faster RCNN模型最近在搞工业质检项目发现Faster R-CNN在缺陷检测上还挺能打。手头有个五类缺陷的数据集划痕、气泡、氧化啥的虽然只有1800张图但配合点数据增强的骚操作效果还不错。今天就跟大伙唠唠实战中那些有意思的细节。先看数据准备这块。VOC格式的好处是能直接用PyTorch的Dataset类不过得处理下XML标注。我习惯用这种写法抓取目标框from xml.etree import ElementTree as ET def parse_voc(xml_path): tree ET.parse(xml_path) boxes [] for obj in tree.findall(object): bndbox obj.find(bndbox) xmin int(bndbox.find(xmin).text) ymin int(bndbox.find(ymin).text) xmax int(bndbox.find(xmax).text) ymax int(bndbox.find(ymax).text) boxes.append([xmin, ymin, xmax, ymax]) return torch.as_tensor(boxes, dtypetorch.float32)注意这里返回的是浮点型张量后面训练时和模型输出的预测框做loss计算才不会报类型错误。遇到过有人在这里用int类型导致训练爆炸的情况排查了半天才发现是数据类型埋的坑。模型搭建部分直接上torchvision的实现省时省力但记得改分类头from torchvision.models.detection import fasterrcnn_resnet50_fpn model fasterrcnn_resnet50_fpn(pretrainedTrue) in_features model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor FastRCNNPredictor(in_features, num_classes6) # 51这里有个新手容易翻车的点——num_classes要算上背景类别。比如我们的缺陷是5类那实际参数要填6。之前有个实习生死活训不出效果最后发现是这里少写了1模型永远在预测背景...基于FasterRCNN目标检测的缺陷检测算法 数据集包含五类别具体如图所示 共计1800张图 包含VOC格式数据集Faster RCNN模型训练策略方面小数据集必备冻结操作。个人喜欢先冻backbone训三天再解冻微调for param in model.backbone.parameters(): param.requires_grad False # 第二阶段解冻 for param in model.backbone[3].parameters(): # 只解冻resnet的最后阶段 param.requires_grad True数据增强这块要结合业务场景。工业缺陷检测最怕几何形变破坏缺陷特征所以用颜色抖动比旋转缩放更靠谱from albumentations import Compose, RGBShift, RandomBrightnessContrast aug Compose([ RGBShift(r_shift_limit15, g_shift_limit15, b_shift_limit15, p0.7), RandomBrightnessContrast(p0.5), ])实测这种组合能让mAP提升3个点左右而且不会像空间变换那样导致漏检。有个有意思的现象气泡缺陷对亮度变化敏感而划痕更吃对比度调整不同类别可能需要差异化的增强策略。最后聊聊推理时的坑。模型输出的是0-1范围的归一化坐标转回原图尺寸时得注意# 假设原图尺寸是(1080, 1920) scale torch.tensor([1920, 1080, 1920, 1080], devicebox.device) restored_box box * scale这里顺序是xmin, ymin, xmax, ymax但scale的张量要对应xyxy的顺序。见过有人把高宽顺序搞反结果框的位置全歪到姥姥家了。实际部署时发现个有趣现象小目标缺陷比如微米级划痕在FPN结构的低层特征图反而检测得更准这可能和深层特征丢失细节有关。后来在RPN阶段调整了anchor的scale设置专门加了组小尺寸anchor召回率立马涨了5%。现在这模型在产线上跑得挺稳平均检测速度在T4显卡上能达到23fps。不过遇到密集小缺陷时还是会有漏检下一步打算把FPN换成更密集的金字塔结构试试。工业场景的实战经验说明没有银弹模型得根据具体缺陷特性反复调教才行。