基于无人机 RGB+红外双模态的小目标行人检测系统与数据集

注：本方案适用于科研/工程落地，可作为 benchmark 系统基础。

一、项目结构

DroneRGBT_Detection/
├── datasets/
│   ├── images/
│   │   ├── train/ # 存放融合后的 4 通道 .npy 或拼接图
│   │   └── val/
│   └── labels/
│       ├── train/
│       └── val/
├── models/
│   └── yolov8s-rgbt.yaml # 自定义 4 通道输入模型配置
├── tools/
│   ├── fuse_rgb_thermal.py # 将 RGB + 红外融合为 4 通道图像
│   └── create_labels.py # （若需从原始标注生成 YOLO 格式）
├── train.py # 训练脚本
├── detect.py # 推理脚本（支持单图/视频）
└── rgbt_drone.yaml # 数据集配置文件

二、融合 RGB + 红外图像

文件：tools/fuse_rgb_thermal.py

import os
import cv2
import numpy as np

def fuse_rgb_thermal(rgb_dir, thermal_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    rgb_files = sorted(os.listdir(rgb_dir))
    for rgb_file in rgb_files:
        if not rgb_file.lower().endswith(('.jpg', '.png')):
            continue
        thermal_file = rgb_file  # 假设文件名一致
        rgb_path = os.path.join(rgb_dir, rgb_file)
        thermal_path = os.path.join(thermal_dir, thermal_file)
        
        rgb = cv2.imread(rgb_path)  # (H, W, 3)
        thermal = cv2.imread(thermal_path, cv2.IMREAD_GRAYSCALE)  # (H, W)
        
        if rgb is None or thermal is None:
            print(f"Skip {rgb_file}")
            continue
        
        # 调整 thermal 到 3 通道（可选）或保留单通道
        thermal = np.expand_dims(thermal, axis=2)  # (H, W, 1)
        fused = np.concatenate([rgb, thermal], axis=2)  # (H, W, 4)
        
        # 保存为 .npy（推荐）或拼接图（如用 4 通道 TIFF）
        np.save(os.path.join(output_dir, rgb_file.replace('.jpg', '.npy')), fused)

if __name__ == '__main__':
    fuse_rgb_thermal(
        rgb_dir='raw_data/train/rgb',
        thermal_dir='raw_data/train/thermal',
        output_dir='datasets/images/train'
    )

✅ 输出为 .npy 文件（含 4 通道），便于后续加载。

三、自定义 YOLOv8 模型（4 通道输入）

创建 models/yolov8s-rgbt.yaml：

# YOLOv8s for RGBT (4-channel input)
nc: 1  # number of classes
scales: [0.33, 0.50]  # model depth and width scaling
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2 (input must be 4-channel!)
  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
  - [-1, 3, C2f, [128, True]]
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  - [-1, 6, C2f, [256, True]]
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]]  # 9
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
  - [-1, 3, C2f, [512]]  # 12
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
  - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 12], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)
  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 9], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)
  - [[15, 18, 21], 1, Detect, [nc]]  # Detect(P3, P4, P5)

⚠️ 关键修改：首层 Conv 输入通道自动适配 4 通道（无需改代码，PyTorch 自动推断）。

四、数据集配置

文件：rgbt_drone.yaml

train: ./datasets/images/train  # 实际存放 .npy 路径
val: ./datasets/images/val
nc: 1
names: ['person']

❗ 注意：Ultralytics 默认读取图像文件（.jpg/.png），但我们需要加载 .npy。因此需自定义数据加载器。

五、自定义数据加载器

在 train.py 中重写 load_image 行为（简化版）：

from ultralytics import YOLO
import torch
import numpy as np
import os
from pathlib import Path
from ultralytics.data.dataset import YOLODataset

# 替换默认的图像加载方式
def custom_load_image(self, i):
    """Load 4-channel .npy image"""
    path = self.files[i]
    if path.endswith('.npy'):
        img = np.load(path)  # (H, W, 4)
        img = img.transpose(2, 0, 1)  # (4, H, W)
        img = torch.from_numpy(img).float()
        return img, img.shape[1], img.shape[2]
    else:
        return self._old_load_image(i)

# Monkey patch（仅用于训练）
YOLODataset._old_load_image = YOLODataset.load_image
YOLODataset.load_image = custom_load_image

# 修改标签加载（确保 .txt 对应 .npy）
def custom_get_label_file(self, img_path):
    return str(Path(img_path).with_suffix('.txt'))

YOLODataset.get_label_file = custom_get_label_file

更稳健做法：继承 YOLODataset 并注册新类，但上述 monkey patch 可快速验证。

六、完整训练脚本

文件：train.py

from ultralytics import YOLO
import torch
import numpy as np
from pathlib import Path
from ultralytics.data.dataset import YOLODataset

# === 自定义数据加载（支持 .npy）===
def load_image_npy(self, i):
    f = self.im_files[i]
    if f.endswith('.npy'):
        im = np.load(f)  # (H, W, 4)
        im = im.transpose(2, 0, 1)  # (4, H, W)
        im = torch.from_numpy(im).float()
        h, w = im.shape[1], im.shape[2]
        return im, h, w
    else:
        return self._orig_load_image(i)

def get_label_file_npy(self, img_path):
    return str(Path(img_path).with_suffix('.txt'))

# Patch
YOLODataset._orig_load_image = YOLODataset.load_image
YOLODataset.load_image = load_image_npy
YOLODataset.get_label_file = get_label_file_npy

# === 开始训练 ===
if __name__ == '__main__':
    model = YOLO('models/yolov8s-rgbt.yaml')  # 从头训练
    # 或加载预训练权重（需修改首层 conv1.weight）
    # model = YOLO('yolov8s.pt')
    # model.model.model[0].conv = torch.nn.Conv2d(4, 64, 3, 2, 1)  # 手动替换
    
    model.train(
        data='rgbt_drone.yaml',
        epochs=100,
        imgsz=512,
        batch=16,
        name='drone_rgbdet',
        project='runs',
        device=0,
        cache=False,
        workers=4
    )

七、推理脚本

文件：detect.py

import torch
import numpy as np
import cv2
from ultralytics import YOLO

model = YOLO('runs/drone_rgbdet/weights/best.pt')

def detect_rgbd_pair(rgb_path, thermal_path):
    rgb = cv2.imread(rgb_path)  # (H, W, 3)
    thermal = cv2.imread(thermal_path, cv2.IMREAD_GRAYSCALE)  # (H, W)
    thermal = np.expand_dims(thermal, axis=2)
    fused = np.concatenate([rgb, thermal], axis=2)  # (H, W, 4)
    fused = fused.transpose(2, 0, 1)  # (4, H, W)
    fused = torch.from_numpy(fused).float().unsqueeze(0) / 255.0  # 归一化
    
    results = model(fused, augment=True)
    
    # 可视化：叠加在 RGB 图上
    annotated = results[0].plot()
    cv2.imshow('Detection', annotated)
    cv2.waitKey(0)

# 示例
detect_rgbd_pair('test/rgb/001.jpg', 'test/thermal/001.jpg')

八、部署建议