基于YOLO的x光安检危险物品检测 数据集介绍:类别为8类,包括:刀(knife)、剪刀(sc...

基于YOLO的x光安检危险物品检测 数据集介绍:类别为8类,包括:刀(knife)、剪刀(scissors)、打火(lighter)、优盘(USBFlashDisk)、压力容器(pressure)、带喷嘴塑料瓶(plasticBottleWithaNozzle)、公章(seal)、电池(battery)。

安检机前的长队总让人焦虑,尤其当X光图像里混着各种奇形怪状的物品时。我们团队最近尝试用YOLOv5解决这个问题时发现,最刺激的不是调参,而是看着模型把公章认成电池的瞬间------这要是真在机场,怕是得引发全员手检了。

先看这个奇葩数据集:8类物品里既有锋利的刀剪,又有看似无害的优盘。特别是带喷嘴的塑料瓶,在X光下会呈现类似手枪握把的轮廓。我们给数据增强上了点狠活:

python 复制代码
train_transforms = A.Compose([
    A.Rotate(limit=15, p=0.6),  # X光包裹经常歪着放
    A.GaussianBlur((3,7), p=0.2),  # 模拟低分辨率设备
    A.RandomBrightnessContrast(0.1, 0.2, p=0.5),
    A.HueSaturationValue(10,15,10,p=0.3),
], bbox_params=A.BboxParams(format='yolo'))

特别注意旋转角度不能超过20度,否则压力容器上的螺纹特征会被扭曲。数据分布也很有意思------打火机的样本量是压力容器的3倍,这逼得我们在损失函数里玩起了平衡术:

python 复制代码
class WeightedLoss(ComputeLoss):
    def __init__(self, model, class_weights):
        super().__init__(model)
        self.class_weights = torch.tensor(class_weights)
        
    def __call__(self, preds, targets):
        loss = super().__call__(preds, targets)
        cls_loss = loss[1] * self.class_weights.mean()
        return loss[0], cls_loss, loss[2]

模型结构选用了YOLOv5s的魔改版,在Backbone里塞了个SPD-Conv模块。这玩意儿特别适合处理X光图像中的锯齿状边缘,原理简单说就是把步幅卷积换成空间到深度的转换:

yaml 复制代码
backbone:
  [[-1, 1, SPDConv, [64, 3]],  # 替换原本的Conv
   [-1, 1, nn.BatchNorm2d, [64]],
   [-1, 1, nn.SiLU, []]]

训练时发现个诡异现象:模型总把横放的剪刀识别为压力容器。最后发现是锚框尺寸不合适,用k-means重新聚类后得到的新锚点:

python 复制代码
新锚框尺寸:
 [12,16, 19,35, 33,23]  # 适合细长物品
 [45,36, 48,65, 79,55]  # 中型容器
 [123,96, 158,183, 324,269]  # 大件物品

推理阶段的骚操作才叫精彩。X光图像往往有金属伪影,我们给NMS加了空间约束------两个重叠框如果出现在图像不同区域(比如包裹的两端),就算IOU超标也保留:

python 复制代码
def nms_with_region(boxes, scores, iou_thresh=0.5, region_thresh=0.3):
    # 根据中心点坐标划分区域
    centers = (boxes[:, :2] + boxes[:, 2:])/2
    region_map = (centers // region_thresh).int()
    # 只在同区域做NMS
    ...

实测发现这个改进让公章和电池的误检率直降15%。现在模型已经能实时处理双视角X光视频流,不过偶尔还是会把卷起来的充电线认成刀具------看来要彻底取代安检员,还得再喂几吨奇葩样本。

相关推荐
皮皮哎哟14 小时前
数据结构:嵌入式常用排序与查找算法精讲
数据结构·算法·排序算法·二分查找·快速排序
堕27415 小时前
java数据结构当中的《排序》(一 )
java·数据结构·排序算法
long31616 小时前
Aho-Corasick 模式搜索算法
java·数据结构·spring boot·后端·算法·排序算法
YuTaoShao19 小时前
【LeetCode 每日一题】3634. 使数组平衡的最少移除数目——(解法一)排序+滑动窗口
算法·leetcode·排序算法
庄周迷蝴蝶1 天前
四、CUDA排序算法实现
算法·排序算法
遨游xyz1 天前
排序-快速排序
开发语言·python·排序算法
历程里程碑2 天前
普通数组----最大子数组和
大数据·算法·elasticsearch·搜索引擎·排序算法·哈希算法·散列表
2302_813806222 天前
【嵌入式修炼:数据结构篇】——单向链表的排序
数据结构·链表·排序算法
酉鬼女又兒2 天前
27. 移除元素
数据结构·算法·排序算法
404未精通的狗2 天前
(数据结构)排序
数据结构·排序算法