基于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, dtype=torch.float32)注意这里返回的是浮点型张量,后面训练时和模型输出的预测框做loss计算才不会报类型错误。遇到过有人在这里用int类型导致训练爆炸的情况,排查了半天才发现是数据类型埋的坑。
模型搭建部分直接上torchvision的实现省时省力,但记得改分类头:
from torchvision.models.detection import fasterrcnn_resnet50_fpn model = fasterrcnn_resnet50_fpn(pretrained=True) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=6) # 5+1这里有个新手容易翻车的点——num_classes要算上背景类别。比如我们的缺陷是5类,那实际参数要填6。之前有个实习生死活训不出效果,最后发现是这里少写了1,模型永远在预测背景...
训练策略方面,小数据集必备冻结操作。个人喜欢先冻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_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.7), RandomBrightnessContrast(p=0.5), ])实测这种组合能让mAP提升3个点左右,而且不会像空间变换那样导致漏检。有个有意思的现象:气泡缺陷对亮度变化敏感,而划痕更吃对比度调整,不同类别可能需要差异化的增强策略。
最后聊聊推理时的坑。模型输出的是0-1范围的归一化坐标,转回原图尺寸时得注意:
# 假设原图尺寸是(1080, 1920) scale = torch.tensor([1920, 1080, 1920, 1080], device=box.device) restored_box = box * scale这里顺序是xmin, ymin, xmax, ymax,但scale的张量要对应xyxy的顺序。见过有人把高宽顺序搞反,结果框的位置全歪到姥姥家了。
实际部署时发现个有趣现象:小目标缺陷(比如微米级划痕)在FPN结构的低层特征图反而检测得更准,这可能和深层特征丢失细节有关。后来在RPN阶段调整了anchor的scale设置,专门加了组小尺寸anchor,召回率立马涨了5%。
现在这模型在产线上跑得挺稳,平均检测速度在T4显卡上能达到23fps。不过遇到密集小缺陷时还是会有漏检,下一步打算把FPN换成更密集的金字塔结构试试。工业场景的实战经验说明,没有银弹模型,得根据具体缺陷特性反复调教才行。