实战分享：用YOLOv5s+小目标检测头搞定红外图像里的‘小不点’（附数据集处理与模型改进）

实战分享用YOLOv5s小目标检测头搞定红外图像里的‘小不点’附数据集处理与模型改进红外图像中的小目标检测一直是计算机视觉领域的难点问题。与常规RGB图像相比红外图像具有低对比度、高噪声等特点这使得传统目标检测方法在红外小目标检测上表现不佳。本文将分享如何基于YOLOv5s模型通过增加小目标检测头和改进数据处理流程有效提升红外图像中小目标的检测精度。1. 红外小目标检测的挑战与解决方案红外图像中的小目标通常指占据图像面积小于32×32像素的物体。这类目标在军事侦察、安防监控、环境监测等领域具有重要应用价值。但红外图像的特殊性给检测带来了三大核心挑战低对比度问题红外图像中目标与背景的温差可能导致边缘模糊高噪声干扰热噪声和传感器噪声会掩盖微小目标的特征多尺度特性同一场景中目标可能呈现从几个像素到上百像素的尺寸变化针对这些挑战我们采用的解决方案框架如下[红外图像输入] → [数据增强处理] → [改进的YOLOv5s网络] → [多尺度预测输出] ↑ ↑ [Mosaic增强] [P2小目标检测头]2. 数据集构建与处理技巧高质量的数据集是模型性能的基础。对于红外小目标检测我们推荐以下数据处理流程2.1 数据采集与标注红外数据集通常有两种获取方式使用公开数据集如infrared_COCO_format自行采集并标注红外图像标注时需特别注意使用YOLO格式的标注工具如LabelImg或YOLO_mark对小目标标注框要尽可能精确减少背景干扰保持标注一致性特别是对模糊目标的判断标准2.2 数据增强策略针对红外小目标的特点我们采用分层增强策略基础增强随机翻转水平/垂直色彩抖动调整亮度、对比度添加高斯噪声高级增强# Mosaic增强示例代码 def mosaic_augmentation(images, labels, size640): 将4张图像拼接为1张 output_image np.zeros((size, size, 3), dtypenp.uint8) output_labels [] # 随机选择拼接点 xc, yc [int(random.uniform(size * 0.25, size * 0.75)) for _ in range(2)] # 处理4个象限 for i in range(4): img, anns random.choice(list(zip(images, labels))) h, w img.shape[:2] if i 0: # 左上 x1a, y1a, x2a, y2a 0, 0, xc, yc x1b, y1b, x2b, y2b 0, 0, w, h elif i 1: # 右上 x1a, y1a, x2a, y2a xc, 0, size, yc x1b, y1b, x2b, y2b 0, 0, w-xc, h elif i 2: # 左下 x1a, y1a, x2a, y2a 0, yc, xc, size x1b, y1b, x2b, y2b 0, 0, w, h-yc else: # 右下 x1a, y1a, x2a, y2a xc, yc, size, size x1b, y1b, x2b, y2b xc, yc, w, h # 裁剪并resize图像块 crop img[y1b:y2b, x1b:x2b] crop cv2.resize(crop, (x2a-x1a, y2a-y1a)) output_image[y1a:y2a, x1a:x2a] crop # 调整标注框坐标 for ann in anns: x, y, w, h ann[bbox] x x1a (x - x1b) * (x2a-x1a)/(x2b-x1b) y y1a (y - y1b) * (y2a-y1a)/(y2b-y1b) w * (x2a-x1a)/(x2b-x1b) h * (y2a-y1a)/(y2b-y1b) output_labels.append([ann[category_id], x, y, w, h]) return output_image, output_labels2.3 数据集组织结构规范的目录结构对训练流程至关重要dataset/ ├── train/ │ ├── images/ # 训练图像 │ └── labels/ # YOLO格式标注文件 ├── val/ │ ├── images/ # 验证图像 │ └── labels/ └── test/ ├── images/ # 测试图像 └── labels/3. YOLOv5s模型改进方案标准YOLOv5s模型对小目标检测效果有限我们通过以下改进提升性能3.1 增加P2小目标检测头YOLOv5默认使用P3、P4、P5三个检测头对应8×8、16×16、32×32的下采样率。我们增加P2检测头4×4下采样专门处理小目标# yolov5s_custom.yaml head: [[-1, 1, Conv, [256, 1, 1]], # 20*20 [-1, 1, nn.Upsample, [None, 2, nearest]], # 40*40 [[-1, 6], 1, Concat, [1]], # cat backbone P4 [-1, 3, BottleneckCSP, [512, False]], # 13 [-1, 1, Conv, [256, 1, 1]], # 40*40 [-1, 1, nn.Upsample, [None, 2, nearest]], # 80*80 [[-1, 4], 1, Concat, [1]], # cat backbone P3 [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) [-1, 1, Conv, [128, 1, 1]], # 18 80*80 [-1, 1, nn.Upsample, [None, 2, nearest]], # 160*160 [[-1, 2], 1, Concat, [1]], # 20 cat backbone p2 [-1, 3, BottleneckCSP, [128, False]], # 21 160*160 (新增P2层) [-1, 1, Conv, [128, 3, 2]], # 22 80*80 [[-1, 18], 1, Concat, [1]], # 23 80*80 [-1, 3, BottleneckCSP, [256, False]], # 24 [-1, 1, Conv, [256, 3, 2]], # 25 40*40 [[-1, 14], 1, Concat, [1]], # 26 cat head P4 [-1, 3, BottleneckCSP, [512, False]], # 27 (P4/16-medium) [-1, 1, Conv, [512, 3, 2]], # 28 20*20 [[-1, 10], 1, Concat, [1]], # 29 cat head P5 [-1, 3, BottleneckCSP, [1024, False]], # 30 (P5/32-large) [[21, 24, 27, 30], 1, Detect, [nc, anchors]], # Detect(P2, P3, P4, P5) ]3.2 自适应锚框计算YOLOv5默认使用COCO数据集的锚框参数对于红外小目标需要重新计算# 计算自定义锚框 python utils/autoanchor.py --cfg yolov5s_custom.yaml --data infrared_data.yaml3.3 训练参数优化针对小目标特点调整训练参数参数常规值小目标优化值说明输入尺寸6401280增大输入保留更多细节batch_size168因显存限制适当减小lr00.010.001小目标需要更精细的学习率mosaic1.00.75避免过度增强导致目标失真fl_gamma0.01.5聚焦困难样本小目标4. 训练技巧与调优经验在实际训练过程中我们总结了以下关键经验4.1 学习率策略采用余弦退火学习率配合线性热身# 学习率调度示例 lr_scheduler torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_05, # 初始周期 T_mult2, # 周期倍增因子 eta_min1e-5 # 最小学习率 )4.2 损失函数调整对小目标检测特别关注分类损失# 修改后的损失计算 class ComputeLoss: def __init__(self, model, autobalanceFalse): self.sort_obj_iou False self.cls_pw 1.0 # 分类损失权重提高 self.obj_pw 1.0 self.fl_gamma 1.5 # Focal Loss参数4.3 验证指标解读重点关注以下指标mAP0.5:0.95综合评估指标mAP0.5常规IoU阈值下的精度mAP_small专门针对小目标的评估结果Recall避免漏检关键目标5. 部署优化与推理加速实际部署时需要考虑效率与精度的平衡5.1 模型量化# 动态量化示例 model_fp32 torch.load(yolov5s_custom.pt) model_int8 torch.quantization.quantize_dynamic( model_fp32, # 原始模型 {torch.nn.Linear, torch.nn.Conv2d}, # 量化层类型 dtypetorch.qint8 # 量化类型 ) torch.save(model_int8, yolov5s_custom_int8.pt)5.2 TensorRT加速# 导出为ONNX python export.py --weights yolov5s_custom.pt --include onnx --img 1280 # 转换为TensorRT trtexec --onnxyolov5s_custom.onnx --saveEngineyolov5s_custom.engine --fp165.3 大图像处理策略对于超高分辨率红外图像如4K可采用分块检测策略将大图像分割为重叠的小块如1280×1280对各小块独立检测合并结果时处理重叠区域的检测框def split_detect(image, model, tile_size1280, overlap0.25): 分块检测大图像 h, w image.shape[:2] preds [] # 计算分块参数 stride int(tile_size * (1 - overlap)) nx (w - tile_size) // stride 1 ny (h - tile_size) // stride 1 # 各分块检测 for i in range(ny1): for j in range(nx1): x1 j * stride y1 i * stride x2 min(x1 tile_size, w) y2 min(y1 tile_size, h) patch image[y1:y2, x1:x2] pred model(patch) # 调整坐标到原图 pred[:, 0] x1 pred[:, 1] y1 preds.append(pred) # 合并结果并NMS all_preds torch.cat(preds, dim0) final_preds non_max_suppression(all_preds) return final_preds6. 实际应用效果与案例在某红外无人机监测项目中改进后的模型表现如下指标原始YOLOv5s改进模型提升幅度mAP0.568.2%82.7%14.5%小目标召回率53.8%76.4%22.6%推理速度(FPS)4532-13模型大小(MB)27.432.14.7典型检测结果对比原始模型漏检远处小目标对密集目标处理不佳改进模型能检测出更多微小目标边界框更精确在实际部署中这套方案已经成功应用于边境安防监控系统电力线路无人机巡检森林防火早期预警红外小目标检测技术的进步为这些应用场景提供了更可靠的技术保障。虽然模型复杂度有所增加但在现代GPU硬件上仍能保持实时性能满足大多数实际应用的需求。