1. 从零开始:认识HIT-UAV数据集与高海拔红外检测的独特挑战

大家好,我是老张,在AI和硬件领域摸爬滚打了十几年,最近几年一直在和无人机、红外图像打交道。今天想和大家深入聊聊一个非常有意思也很有挑战性的实战项目:如何用YOLOv5搞定HIT-UAV这个高海拔红外数据集的目标检测。如果你对无人机巡检、安防监控或者恶劣环境下的视觉任务感兴趣,这篇内容应该能给你不少启发。

首先,我们得搞清楚HIT-UAV是个啥。简单说,它就是哈尔滨工业大学专门为无人机在高海拔地区执行任务而创建的一个红外热成像数据集。想象一下,在高原、雪山或者边境地带,普通摄像头可能因为光照、雾气、沙尘而“失明”,但红外摄像头靠感知物体散发的热量成像,基本不受这些影响。HIT-UAV里面全是无人机从高空拍下来的红外图片,目标包括车辆、小型建筑、甚至是一些自然物体,每张图都精细标注了目标的位置框。

但为什么说它“有挑战”呢?我实测下来,高海拔红外图像和咱们平时玩的COCO、VOC这些可见光数据集差别太大了。第一是目标特性模糊。红外图像没有纹理和颜色,全靠温差形成的亮度差异来区分物体。一辆车和一块被太阳晒热的石头,在图像上可能都是个亮斑,边界非常不清晰。第二是尺度变化剧烈。无人机飞得高,地面目标看起来就是一个个小点,而且由于透视关系,同一类目标在图像边缘和中心的大小差异巨大。第三是背景复杂。高原环境地形地貌多变,山体阴影、云层热辐射都会形成干扰,很容易让模型产生误报。

所以,直接把现成的YOLOv5模型丢上去训练,效果通常不会太好。我们需要一个完整的流程,从理解数据开始,到针对性地改造模型和训练策略。这个过程就像给一位习惯了城市道路的司机做特训,让他能在青藏高原上安全驾驶一样,需要调整的细节非常多。接下来,我就带你一步步走通这个全流程,分享我踩过的坑和验证有效的技巧。

2. 数据到手第一步:HIT-UAV的获取、解析与“体检”

拿到数据集是万里长征的第一步。HIT-UAV通常需要向官方申请获取,这里假设你已经拿到了数据压缩包。解压后,别急着跑代码,我们先给数据做个全面的“体检”。

2.1 深入理解数据结构与内容

解压后的文件夹结构可能如下,但不同版本略有差异,核心是找到图片和标注的对应关系:

HIT-UAV/
├── images/
│   ├── seq_001_frame_000001.jpg
│   ├── seq_001_frame_000002.jpg
│   └── ...
├── annotations/
│   ├── seq_001_frame_000001.json
│   ├── seq_001_frame_000002.json
│   └── ...
└── dataset_info.json

和常见数据集不同,HIT-UAV的标注通常是JSON格式,里面信息很丰富。我们写段代码打开一个看看:

import json
import os

# 加载一个标注文件示例
annotation_path = 'HIT-UAV/annotations/seq_001_frame_000001.json'
with open(annotation_path, 'r') as f:
    ann = json.load(f)

print(f"图像文件名: {ann['image']['file_name']}")
print(f"图像尺寸: 宽{ann['image']['width']}, 高{ann['image']['height']}")
print(f"标注目标数量: {len(ann['annotations'])}")

# 查看第一个目标的详细信息
first_obj = ann['annotations'][0]
print(f"\n目标类别: {first_obj['category_name']} (ID: {first_obj['category_id']})")
print(f"边界框坐标 [x_min, y_min, 宽, 高]: {first_obj['bbox']}")
print(f"边界框面积: {first_obj['area']}")
print(f"是否被遮挡: {first_obj['occluded']}")
print(f"是否为小目标: {first_obj['is_small']}")

运行后你可能会发现,标注里不仅给了边界框[x_min, y_min, width, height],还有occluded(是否被遮挡)、is_small(是否是小目标)这样的属性标签,这对我们后续分析数据难点、设计数据增强策略至关重要。比如,如果is_small=True的比例很高,那我们就知道小目标检测是主要矛盾。

2.2 数据可视化与统计分析:发现“病灶”

光看数字不够直观,我们必须把图片和标注框画出来看看。同时,做一些统计分析,做到心中有数。

from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

def visualize_sample(image_id):
    """可视化单张图片及其标注"""
    img_path = os.path.join('HIT-UAV/images', f"{image_id}.jpg")
    ann_path = os.path.join('HIT-UAV/annotations', f"{image_id}.json")

    img = Image.open(img_path).convert('L')  # 红外图是单通道灰度图
    with open(ann_path, 'r') as f:
        ann = json.load(f)

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    axes[0].imshow(img, cmap='hot')  # 使用‘hot’色彩映射更接近热成像视觉效果
    axes[0].set_title('原始红外图像 (Hot Colormap)')
    axes[0].axis('off')

    axes[1].imshow(img, cmap='gray')
    for obj in ann['annotations']:
        bbox = obj['bbox']
        # 将 [x_min, y_min, width, height] 转换为矩形绘制所需的坐标
        rect = plt.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3],
                             linewidth=1.5, edgecolor='lime', facecolor='none')
        axes[1].add_patch(rect)
        # 在框上方显示类别
        axes[1].text(bbox[0], bbox[1]-5, obj['category_name'],
                     color='lime', fontsize=9, weight='bold')
    axes[1].set_title('标注可视化')
    axes[1].axis('off')
    plt.tight_layout()
    plt.show()

# 随机看几张
sample_ids = ['seq_001_frame_000010', 'seq_002_frame_000050', 'seq_003_frame_000100']
for sid in sample_ids:
    visualize_sample(sid)

通过可视化,你能立刻感受到我之前说的:目标对比度低、边缘模糊、尺寸小。接下来,我们做点统计,用数据说话:

def dataset_statistics(annotation_dir):
    """统计数据集关键信息"""
    ann_files = [f for f in os.listdir(annotation_dir) if f.endswith('.json')]
    total_objs = 0
    size_distribution = {'tiny':0, 'small':0, 'medium':0, 'large':0}
    category_count = {}
    occlusion_count = 0

    for ann_file in ann_files[:500]:  # 抽样统计,加快速度
        with open(os.path.join(annotation_dir, ann_file), 'r') as f:
            ann = json.load(f)
        for obj in ann['annotations']:
            total_objs += 1
            # 统计类别
            cat = obj['category_name']
            category_count[cat] = category_count.get(cat, 0) + 1
            # 根据面积判断目标大小 (这里按像素面积粗略划分)
            area = obj['area']
            if area < 32*32:
                size_distribution['tiny'] += 1
            elif area < 96*96:
                size_distribution['small'] += 1
            elif area < 256*256:
                size_distribution['medium'] += 1
            else:
                size_distribution['large'] += 1
            # 统计遮挡
            if obj.get('occluded', False):
                occlusion_count += 1

    print(f"统计样本数: {len(ann_files)}")
    print(f"总目标实例数: {total_objs}")
    print(f"\n目标尺寸分布:")
    for k, v in size_distribution.items():
        print(f"  {k}: {v} ({v/total_objs*100:.1f}%)")
    print(f"\n目标类别分布:")
    for cat, cnt in category_count.items():
        print(f"  {cat}: {cnt}")
    print(f"\n被遮挡目标比例: {occlusion_count/total_objs*100:.1f}%")

dataset_statistics('HIT-UAV/annotations')

跑完这个统计,你很可能发现tiny和small目标占了绝大多数(可能超过70%),这就是高海拔无人机视角的典型特征。这个“体检报告”直接决定了我们后续所有技术方案的选择:模型必须对小目标友好,数据增强要能提升小目标的鲁棒性,损失函数也要更关注小目标。

3. 为YOLOv5定制数据预处理与增强流水线

数据看明白了,下一步就是为YOLOv5准备“食粮”。直接喂原始数据效果差,我们必须根据HIT-UAV的特点,设计一套预处理和增强组合拳。

3.1 标注格式转换:从JSON到YOLO TXT

YOLOv5训练需要特定的TXT格式,每行代表一个目标:class_id x_center y_center width height,所有坐标都是相对于图片宽高的归一化值(0到1之间)。转换时有个关键点:HIT-UAV的JSON标注里bbox是[x_min, y_min, width, height],而YOLO格式需要中心点坐标和宽高。

import os
import json
from tqdm import tqdm

def convert_hit_uav_to_yolo(json_dir, output_label_dir, class_map):
    """
    将HIT-UAV的JSON标注转换为YOLOv5所需的TXT格式。
    Args:
        json_dir: JSON标注文件夹路径
        output_label_dir: 输出TXT标签文件夹路径
        class_map: 类别名称到ID的映射字典,如 {'vehicle': 0, 'building': 1}
    """
    os.makedirs(output_label_dir, exist_ok=True)
    json_files = [f for f in os.listdir(json_dir) if f.endswith('.json')]

    for json_file in tqdm(json_files, desc="转换标注格式"):
        json_path = os.path.join(json_dir, json_file)
        with open(json_path, 'r') as f:
            data = json.load(f)

        img_width = data['image']['width']
        img_height = data['image']['height']
        txt_lines = []

        for obj in data['annotations']:
            cat_name = obj['category_name']
            if cat_name not in class_map:
                continue  # 跳过未定义类别
            class_id = class_map[cat_name]

            # 获取bbox: [x_min, y_min, width, height]
            x_min, y_min, w, h = obj['bbox']
            # 计算中心点坐标并归一化
            x_center = (x_min + w / 2.0) / img_width
            y_center = (y_min + h / 2.0) / img_height
            width_norm = w / img_width
            height_norm = h / img_height

            # 确保坐标在[0,1]范围内(处理可能的标注误差)
            x_center = max(0, min(1, x_center))
            y_center = max(0, min(1, y_center))
            width_norm = max(0, min(1, width_norm))
            height_norm = max(0, min(1, height_norm))

            txt_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width_norm:.6f} {height_norm:.6f}")

        # 写入TXT文件,文件名与图片名对应
        txt_filename = os.path.splitext(json_file)[0] + '.txt'
        txt_path = os.path.join(output_label_dir, txt_filename)
        with open(txt_path, 'w') as f:
            f.write('\n'.join(txt_lines))

# 定义类别映射(根据你的数据集实际类别调整)
class_mapping = {
    'vehicle': 0,
    'building': 1,
    'person': 2,
    # ... 添加其他类别
}
convert_hit_uav_to_yolo('HIT-UAV/annotations', 'HIT-UAV/labels', class_mapping)

转换完成后,记得按YOLOv5的要求组织数据集文件夹,通常是这样:

HIT-UAV_YOLO/
├── images/
│   ├── train/
│   └── val/
└── labels/
    ├── train/
    └── val/

你需要自己写个脚本,按一定比例(如8:2)把图片和对应的TXT标签文件分别移动到train和val文件夹下。

3.2 针对红外与小目标的增强策略设计

这是提升模型性能的关键一步。YOLOv5自带的增强(如mosaic, mixup)很强,但针对红外小目标,我们需要更有针对性的调整。我建议修改YOLOv5的data/hyps/hyp.scratch-low.yaml配置文件(如果用预训练模型,则用hyp.finetune.yaml),或者直接在训练命令里覆盖相关参数。

首先,增大输入分辨率。高海拔小目标在640x640的输入下可能只有几个像素,网络根本学不到特征。我尝试把--img参数调到1024甚至1280,对小目标召回率提升明显,但显存消耗和速度会下降,需要权衡。

python train.py --img 1024 --batch-size 8 --epochs 100 ... # 根据你的GPU调整batch-size

其次,调整增强参数。在hyp配置文件中,我通常会做这些修改:

# hyp.custom_hit_uav.yaml
hsv_h: 0.015  # 色调增强减弱(红外图色调信息少)
hsv_s: 0.7    # 饱和度增强加强(模拟不同热辐射强度)
hsv_v: 0.4    # 明度增强适中
degrees: 5.0  # 旋转角度减小,避免小目标转出视野
translate: 0.1  # 平移减弱
scale: 0.9   # 缩放增强减弱,避免目标缩得太小
shear: 2.0    # 剪切减弱
perspective: 0.0005  # 透视变换减弱
flipud: 0.0   # 上下翻转概率设为0(无人机视角一般不会上下颠倒)
fliplr: 0.5   # 左右翻转保留
mosaic: 1.0   # Mosaic增强保持开启,对丰富小目标上下文很有用
mixup: 0.1    # Mixup比例调低,避免过度模糊小目标特征
copy_paste: 0.3 # 开启并调整copy-paste增强,这对增加小目标样本数量非常有效!

这里重点说下copy_paste,这是YOLOv5后期版本加入的增强,它会随机复制一些目标粘贴到图片的其他位置。对于HIT-UAV这种小目标多的数据集,它能显著增加正样本数量,我实测下来对AP_s(小目标平均精度)提升能有3-5个百分点。

最后,可以自定义增强。比如,红外图像对高斯噪声和模拟热噪声更敏感,我们可以在datasets.py里添加自定义的噪声增强层。也可以在训练前对全数据集做一次自适应直方图均衡化(CLAHE),来增强低对比度目标的边缘,这个预处理步骤对红外图像效果拔群。

# 示例:在加载图像后,应用CLAHE增强(可集成到数据加载器中)
import cv2
def apply_clahe(image):
    """对单通道红外图像进行CLAHE增强"""
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    enhanced = clahe.apply(image)
    return enhanced

4. YOLOv5模型的高海拔红外适配改造

数据准备好了,模型本身也需要“特调”。直接用原始的YOLOv5s/m/l/x,可能不是最优解。我们需要针对“红外”、“小目标”、“高海拔俯视”这三个关键词进行模型层面的微调。

4.1 注意力机制引入:让模型聚焦关键热源

红外图像中,目标与背景的区分主要靠温差(亮度差)。传统的卷积层对所有区域一视同仁,我们可以加入注意力机制,让网络学会重点关注那些“亮”的、可能是热源的区域。一个简单有效的方法是在Backbone(特征提取网络)的末端加入CBAM(Convolutional Block Attention Module) 或 SE(Squeeze-and-Excitation) 模块。

以在YOLOv5的SPPF层之前加入CBAM为例,我们需要修改models/common.py文件:

# 在 common.py 中添加CBAM模块定义
import torch
import torch.nn as nn

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        out = avg_out + max_out
        return self.sigmoid(out)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x_cat = torch.cat([avg_out, max_out], dim=1)
        out = self.conv(x_cat)
        return self.sigmoid(out)

class CBAM(nn.Module):
    def __init__(self, c1, ratio=16, kernel_size=7):
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttention(c1, ratio)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        out = x * self.channel_attention(x)
        out = out * self.spatial_attention(out)
        return out

# 然后,在 models/yolo.py 中修改你的模型解析逻辑,在需要的地方插入CBAM模块。
# 更简单的方法:直接修改YOLOv5的模型配置文件(.yaml文件)。

你可以创建一个新的模型配置文件,例如yolov5s_cbam.yaml,在Backbone的最后一层(通常是第9层)后面插入CBAM模块。这样,模型在提取高层语义特征时,会同时考虑通道和空间上的重要性,对热源目标更敏感。

4.2 特征融合网络优化:提升小目标检测能力

YOLOv5的Neck部分(FPN+PAN)负责融合不同尺度的特征。但对于大量像素点大小的小目标,深层的特征图(如80x80)经过多次下采样,小目标的信息可能已经丢失。为此,我们可以:

  1. 增加更浅层的特征输出:将Backbone中较浅层(例如C3层)的特征也引入到Neck中进行融合。这需要修改Neck结构,增加额外的上采样和融合路径,相当于构建一个“高分辨率”的特征金字塔,让小目标在融合时仍有较丰富的细节。
  2. 使用更高效的上下采样模块:将普通的卷积下采样替换为Focus模块或跨阶段局部网络(CSPNet) 的变体,减少信息损失。YOLOv5本身已部分采用,我们可以检查其配置。
  3. 引入即插即用的轻量级模块:如BiFPN(加权双向特征金字塔网络),它通过可学习的权重来融合不同尺度的特征,强调对小目标更重要的特征层。在YOLOv5中集成BiFPN需要对Neck部分进行重写,工程量较大,但效果提升也显著,尤其在小目标密集的场景。

一个更务实的做法是直接使用YOLOv5-P2版本(如果有)或寻找社区修改版,它们通常扩展了P2(160x160尺度)的输出层,专门用于检测微小目标。如果自己改,思路是在Head部分增加一个针对更大特征图(下采样倍数小,如4倍)的检测头。

4.3 损失函数调优:平衡大小目标的贡献

默认的CIoU Loss对于大小目标一视同仁。但在HIT-UAV中,小目标数量占优,大目标虽然少但同样重要。我们可以尝试:

  • 使用WIoU(Wise-IoU):最新版本的YOLOv5/v8已经支持。WIoU通过动态调整损失权重,减轻简单大样本对梯度的主导,让模型更关注难例(其中很多是小目标)。
  • 调整损失权重:修改utils/loss.py中的box_loss_gain、cls_loss_gain和obj_loss_gain。对于小目标多的场景,可以适当提高cls_loss_gain(分类损失权重),因为小目标分类更难;同时,也可以微调box_loss_gain。但这需要谨慎的消融实验。
  • 引入Focal Loss的思想:虽然YOLOv5的obj loss本身有一定focal效果,但对于极端类别不平衡(背景远多于目标),可以在分类损失中尝试Focal Loss来压制大量简单负样本(背景)的贡献。

在训练命令中,可以指定使用WIoU:

python train.py ... --bbox_wiou  # 如果版本支持

或者,更直接的方法是修改hyp配置文件中的box损失权重参数,并观察验证集上AP_s和AP_m、AP_l的变化,找到一个平衡点。

5. 训练、调参与性能优化实战

一切就绪,开始训练。这里不是简单地跑起来就行,而是要根据训练过程中的反馈,进行精细化的调参和优化。

5.1 训练启动与关键监控指标

假设我们已经按YOLOv5要求组织好数据,并创建了数据集配置文件HIT-UAV.yaml:

# HIT-UAV.yaml
path: /path/to/HIT-UAV_YOLO  # 数据集根目录
train: images/train  # 训练集路径(相对path)
val: images/val      # 验证集路径

# 类别数
nc: 3
# 类别名称列表,必须与转换标注时的class_mapping顺序一致
names: ['vehicle', 'building', 'person']

启动训练的命令可能看起来像这样:

cd /path/to/yolov5
python train.py \
  --img 1024 \
  --batch-size 8 \
  --epochs 150 \
  --data ./data/HIT-UAV.yaml \
  --cfg ./models/yolov5s_cbam.yaml \  # 使用我们自定义的带CBAM的模型配置
  --weights yolov5s.pt \
  --hyp ./data/hyps/hyp.custom_hit_uav.yaml \  # 使用自定义的超参数
  --name hit_uav_exp1 \
  --project HIT-UAV_runs \
  --cache ram  # 使用RAM缓存加速数据加载(如果内存足够大)

训练开始后,重点监控TensorBoard或W&B上的这些曲线:

  1. 损失曲线:train/box_loss, train/obj_loss, train/cls_loss 以及对应的val损失。理想情况是它们平稳下降,且训练和验证损失没有巨大差距(防过拟合)。
  2. 性能指标:metrics/mAP_0.5, metrics/mAP_0.5:0.95。更重要的是**metrics/precision和metrics/recall**。高海拔红外检测往往召回率(Recall)偏低(很多小目标漏检),我们的优化要朝着提升Recall的方向努力。
  3. 特定尺寸AP:如果日志能输出AP_small, AP_medium, AP_large,那太好了,直接盯着AP_small看。

5.2 超参数调优与技巧分享

根据我多次实验的经验,有几个超参数对HIT-UAV影响巨大:

  • 学习率(lr0):红外特征与ImageNet预训练模型的可见光特征分布不同,建议使用稍大的初始学习率(如0.01),并配合cos或linear的学习率调度器,让模型更快地适应新数据分布。可以在hyp文件中设置lr0: 0.01,lrf: 0.01(最终学习率为初始的1%)。
  • 优化器选择:尝试使用AdamW代替默认的SGD。AdamW对于这种数据分布差异大的任务,有时能更快收敛到更好的位置。在训练命令中加入--optimizer AdamW。
  • 热身(Warmup):务必开启。--warmup-epochs 3和--warmup-momentum 0.8可以帮助模型在初始阶段稳定地探索参数空间,避免早期震荡。
  • 多尺度训练:--multi-scale 参数在YOLOv5中默认是关闭的。对于高海拔数据集,开启多尺度训练(如--multi-scale --img-size 640 1024)能让模型适应不同飞行高度(对应不同目标尺度)的成像,显著提升泛化能力,但会延长训练时间。
  • 早停(Early Stopping):耐心值--patience可以设为50或更高。因为红外小目标检测的训练过程可能会有较长的平台期,不要过早停止。

5.3 模型验证与结果分析

训练完成后,用最佳权重在验证集上做全面评估:

python val.py \
  --weights ./HIT-UAV_runs/hit_uav_exp1/weights/best.pt \
  --data ./data/HIT-UAV.yaml \
  --img 1024 \
  --batch-size 16 \
  --task val \
  --verbose \
  --save-json  # 输出详细的JSON结果文件,便于分析

分析结果时,不要只看一个mAP。打开生成的results.json或查看终端输出,重点关注:

  • 每个类别的AP:看看是不是某个类别(比如person)特别难检测。
  • 混淆矩阵:运行python val.py ... --save-confusion-matrix,生成混淆矩阵图。这能告诉你模型最容易把哪些类别搞混(例如,把vehicle误检为building),从而指导你后续的数据清洗或类别合并。
  • PR曲线:特别是每个类别的PR曲线。如果某个类别的曲线靠近坐标轴,说明该类检测效果很差。
  • 可视化预测:用detect.py脚本在验证集上跑一些图片,直观感受问题所在。
python detect.py \
  --weights ./HIT-UAV_runs/hit_uav_exp1/weights/best.pt \
  --source ./HIT-UAV_YOLO/images/val/ \
  --conf-thres 0.25 \
  --iou-thres 0.45 \
  --save-txt \
  --save-conf \
  --exist-ok

打开runs/detect/exp文件夹,查看预测结果。你会清晰地看到哪些小目标被漏掉了,哪些背景区域被误检了。这个直观反馈是下一步迭代优化最重要的依据。可能是某个区域的热噪声被当成了目标,也可能是远处密集的小车辆被模型合并成了一个框。针对这些问题,再回头去调整数据增强、模型结构或损失函数,进入下一个优化循环。这个过程可能重复多次,但每一次你都会对数据和模型有更深的理解,效果也会逐步提升。

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐