HIT-UAV：无人机高海拔红外热成像目标检测数据集

无人机（UAV）应用在多个领域，尤其是在高海拔环境中，成为了一个非常有价值的技术。HIT-UAV 是一个专门为无人机高海拔红外热成像目标检测任务提供的数据集，它包含了高海拔地区拍摄的红外热成像数据，能够用于训练和评估目标检测模型。本篇文章将介绍 HIT-UAV 数据集，如何加载、处理并应用该数据集进行目标检测任务，并提供丰富的代码示例。

---

一、HIT-UAV 数据集简介

HIT-UAV 数据集由哈尔滨工业大学（HIT）提供，旨在为高海拔环境下的红外热成像目标检测任务提供高质量的标注数据。数据集包含了大量来自无人机在高海拔地区飞行时拍摄的红外图像，每张图像都包含了多种目标（如车辆、建筑物、树木等）以及相应的标注信息。

• 图像类型：红外热成像图像
• 数据来源：高海拔地区的无人机飞行
• 标注类型：目标边界框（bounding box）
• 目标类型：车辆、建筑物、树木等
• 图像分辨率：高分辨率的红外热成像图像（例如 1024x1024 像素）

该数据集主要用于训练和评估目标检测算法，尤其是那些在复杂环境中进行目标检测的算法。

---

二、数据集获取与加载

在开始进行模型训练和评估之前，首先需要获取 HIT-UAV 数据集。你可以通过以下步骤下载和加载数据集：

1. 数据集下载

目前，HIT-UAV 数据集可以通过向提供方申请获取。你需要联系数据集的提供方，并按照他们的要求完成注册和下载。一般来说，数据集会以压缩文件的形式提供，包含图像和标注信息。

2. 数据集结构

假设我们已经下载并解压了数据集，数据集的基本结构可能如下：

HIT-UAV/
│
├── images/       # 红外图像文件夹
│   ├── img_001.jpg
│   ├── img_002.jpg
│   └── ...
│
├── annotations/  # 标注文件夹
│   ├── img_001.json
│   ├── img_002.json
│   └── ...
│
└── README.md

• images/：存储红外热成像图像
• annotations/：存储对应图像的标注文件，通常为 JSON 格式，包含了目标的位置（边界框）信息。

3. 加载数据集

为了加载和处理这些数据，我们可以编写一个 Python 脚本，使用 Pillow 来加载图像，使用 json 来解析标注文件。以下是一个示例代码，展示如何加载图像及其标注信息：

import os
import json
from PIL import Image
import matplotlib.pyplot as plt

# 定义数据集路径
image_dir = 'HIT-UAV/images/'
annotation_dir = 'HIT-UAV/annotations/'

# 加载图像和标注
def load_data(image_id):
    # 加载图像
    image_path = os.path.join(image_dir, f'{image_id}.jpg')
    image = Image.open(image_path)

    # 加载标注
    annotation_path = os.path.join(annotation_dir, f'{image_id}.json')
    with open(annotation_path, 'r') as f:
        annotations = json.load(f)

    return image, annotations

# 可视化图像及其标注
def display_image_with_annotations(image, annotations):
    plt.imshow(image)
    for annotation in annotations['objects']:
        x1, y1, x2, y2 = annotation['bbox']
        plt.gca().add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, color='red', linewidth=2))
    plt.show()

# 示例加载和展示第一张图像及其标注
image_id = 'img_001'  # 图像 ID
image, annotations = load_data(image_id)
display_image_with_annotations(image, annotations)

三、目标检测模型训练

加载数据集并进行预处理后，接下来是训练目标检测模型。这里我们使用 YOLOv5（一个流行的目标检测算法）进行训练。YOLOv5 是一个高效的目标检测模型，适用于各种环境下的实时检测。

1. 安装依赖

在开始之前，确保安装了必要的依赖：

pip install torch torchvision
pip install yolov5

2. 数据格式转换

YOLOv5 需要的数据格式为 .txt 文件，每个 .txt 文件对应一张图像，包含每个目标的类别、位置（边界框坐标）等信息。我们需要将 HIT-UAV 数据集的标注转换为 YOLO 格式。

以下是一个将 HIT-UAV 数据集的标注文件转换为 YOLO 格式的代码示例：

import os
import json

# 转换标注为 YOLO 格式
def convert_to_yolo_format(image_id, annotations):
    # 获取图像的宽高
    image_width, image_height = 1024, 1024  # 假设图像为 1024x1024
    
    # 获取目标类别和边界框
    yolo_annotations = []
    for obj in annotations['objects']:
        category = obj['category']
        x1, y1, x2, y2 = obj['bbox']

        # 转换为 YOLO 格式 [class_id, x_center, y_center, width, height]
        x_center = (x1 + x2) / 2 / image_width
        y_center = (y1 + y2) / 2 / image_height
        width = (x2 - x1) / image_width
        height = (y2 - y1) / image_height
        
        # 假设类别映射为 0, 1, 2, ..., 你可以根据实际情况调整
        class_id = 0  # 示例，假设只有一个类别
        
        yolo_annotations.append(f'{class_id} {x_center} {y_center} {width} {height}')
    
    # 保存为 YOLO 格式文件
    output_path = os.path.join('yolo_labels', f'{image_id}.txt')
    with open(output_path, 'w') as f:
        f.write('\n'.join(yolo_annotations))

# 示例转换
image_id = 'img_001'
image, annotations = load_data(image_id)
convert_to_yolo_format(image_id, annotations)

3. 训练 YOLOv5 模型

一旦数据集准备好并转换为 YOLO 格式，就可以使用 YOLOv5 进行训练。以下是训练 YOLOv5 模型的基本步骤：

1. 克隆 YOLOv5 的 GitHub 仓库：

   git clone https://github.com/ultralytics/yolov5
   cd yolov5
   

2. 将数据集放入 YOLOv5 所需的文件夹结构中：

   yolov5/
   ├── data/
   │   ├── HIT-UAV.yaml  # 数据集配置文件
   │   ├── images/
   │   └── labels/
   ├── models/
   └── ...
   

3. 创建数据集配置文件 HIT-UAV.yaml，内容类似如下：

   train: ../data/images/train
   val: ../data/images/val
   
   nc: 1  # 类别数
   names: ['target']  # 类别名称
   

4. 训练模型：

   python train.py --img 640 --batch 16 --epochs 50 --data data/HIT-UAV.yaml --weights yolov5s.pt
   

4. 评估与测试

训练完成后，你可以使用以下命令进行评估和测试：

python val.py --weights runs/train/exp/weights/best.pt --data data/HIT-UAV.yaml --img 640

---

四、总结

HIT-UAV 数据集为高海拔无人机红外热成像目标检测任务提供了丰富的数据资源。通过合适的数据加载、预处理以及目标检测模型训练（如 YOLOv5），可以有效地进行目标检测。本文介绍了如何从数据集获取、加载、转换数据格式到训练模型的完整流程，并提供了相应的代码示例。通过这些步骤，你可以使用 HIT-UAV 数据集进行高海拔环境下的目标检测任务研究。