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

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

红外图像中的小目标检测一直是计算机视觉领域的难点问题。与常规RGB图像相比，红外图像具有低对比度、高噪声等特点，这使得传统目标检测方法在红外小目标检测上表现不佳。本文将分享如何基于YOLOv5s模型，通过增加小目标检测头和改进数据处理流程，有效提升红外图像中小目标的检测精度。

1. 红外小目标检测的挑战与解决方案

红外图像中的小目标通常指占据图像面积小于32×32像素的物体。这类目标在军事侦察、安防监控、环境监测等领域具有重要应用价值。但红外图像的特殊性给检测带来了三大核心挑战：

1. 低对比度问题：红外图像中目标与背景的温差可能导致边缘模糊
2. 高噪声干扰：热噪声和传感器噪声会掩盖微小目标的特征
3. 多尺度特性：同一场景中目标可能呈现从几个像素到上百像素的尺寸变化

针对这些挑战，我们采用的解决方案框架如下：

[红外图像输入] → [数据增强处理] → [改进的YOLOv5s网络] → [多尺度预测输出] ↑ ↑ [Mosaic增强] [P2小目标检测头]

2. 数据集构建与处理技巧

高质量的数据集是模型性能的基础。对于红外小目标检测，我们推荐以下数据处理流程：

2.1 数据采集与标注

红外数据集通常有两种获取方式：

• 使用公开数据集（如infrared_COCO_format）
• 自行采集并标注红外图像

标注时需特别注意：

• 使用YOLO格式的标注工具（如LabelImg或YOLO_mark）
• 对小目标标注框要尽可能精确，减少背景干扰
• 保持标注一致性，特别是对模糊目标的判断标准

2.2 数据增强策略

针对红外小目标的特点，我们采用分层增强策略：

基础增强：

• 随机翻转（水平/垂直）
• 色彩抖动（调整亮度、对比度）
• 添加高斯噪声

高级增强：

# Mosaic增强示例代码 def mosaic_augmentation(images, labels, size=640): """ 将4张图像拼接为1张 """ output_image = np.zeros((size, size, 3), dtype=np.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_labels

2.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.yaml

3.3 训练参数优化

针对小目标特点调整训练参数：

4. 训练技巧与调优经验

在实际训练过程中，我们总结了以下关键经验：

4.1 学习率策略

采用余弦退火学习率配合线性热身：

# 学习率调度示例 lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=5, # 初始周期 T_mult=2, # 周期倍增因子 eta_min=1e-5 # 最小学习率 )

4.2 损失函数调整

对小目标检测特别关注分类损失：

# 修改后的损失计算 class ComputeLoss: def __init__(self, model, autobalance=False): self.sort_obj_iou = False self.cls_pw = 1.0 # 分类损失权重提高 self.obj_pw = 1.0 self.fl_gamma = 1.5 # Focal Loss参数

4.3 验证指标解读

重点关注以下指标：

• mAP@0.5:0.95：综合评估指标
• mAP@0.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}, # 量化层类型 dtype=torch.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 --onnx=yolov5s_custom.onnx --saveEngine=yolov5s_custom.engine --fp16

5.3 大图像处理策略

对于超高分辨率红外图像（如4K），可采用分块检测策略：

1. 将大图像分割为重叠的小块（如1280×1280）
2. 对各小块独立检测
3. 合并结果时处理重叠区域的检测框

def split_detect(image, model, tile_size=1280, overlap=0.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(ny+1): for j in range(nx+1): 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, dim=0) final_preds = non_max_suppression(all_preds) return final_preds

6. 实际应用效果与案例

在某红外无人机监测项目中，改进后的模型表现如下：

典型检测结果对比：

• 原始模型：漏检远处小目标，对密集目标处理不佳
• 改进模型：能检测出更多微小目标，边界框更精确

在实际部署中，这套方案已经成功应用于：

• 边境安防监控系统
• 电力线路无人机巡检
• 森林防火早期预警

红外小目标检测技术的进步为这些应用场景提供了更可靠的技术保障。虽然模型复杂度有所增加，但在现代GPU硬件上仍能保持实时性能，满足大多数实际应用的需求。