基于anti-uav数据集的目标检测模型优化实战：从数据清洗到模型部署

1次阅读

没有评论

共计 2695 个字符，预计需要花费 7 分钟才能阅读完成。

Anti-UAV 数据集是专门用于无人机检测的公开数据集，但在实际使用过程中，我们遇到了几个典型的挑战：

小目标占比高：无人机在远距离拍摄时往往只占几个像素，容易被传统检测算法忽略
运动模糊：快速移动的无人机容易产生运动模糊，影响特征提取
类别不平衡：正样本（无人机）远少于负样本（背景）
复杂背景干扰：天空云层、建筑边缘等容易产生误检

我们对比了几种主流目标检测模型在 anti-uav 数据集上的表现：

Faster R-CNN：mAP@0.5 达到 68.2%，但推理速度仅 15FPS（NVIDIA T4）
YOLOv5s：速度达到 45FPS，但小目标召回率只有 62.3%
YOLOv5x：mAP@0.5 提升至 72.1%，速度降至 28FPS

最终选择 YOLOv5x 作为基础模型，因为：

单阶段检测器速度优势明显
通过调整 FPN 结构可以针对性优化小目标检测
社区支持完善，便于部署

# 自定义数据加载器
class UAVDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = Path(img_dir)
        self.transform = transform

    def __getitem__(self, idx):
        img_path = self.img_dir / f'{idx}.jpg'
        label_path = img_path.with_suffix('.txt')

        img = cv2.imread(str(img_path))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 处理运动模糊
        if random.random() > 0.7:
            img = cv2.GaussianBlur(img, (3,3), 0)

        # 小目标增强
        if random.random() > 0.5:
            img = cv2.resize(img, (0,0), fx=1.2, fy=1.2)

        # 其余 transform...
        return img, labels

调整 FPN 结构：在 P2 层（1/ 4 尺度）增加检测头
修改 anchor 尺寸：使用 k -means 重新聚类生成适合小目标的 anchor
添加注意力模块：在 backbone 末端插入 CBAM 模块

# YOLOv5 模型修改示例
class Detect(nn.Module):
    def __init__(self, nc=80, anchors=()):
        super().__init__()
        # 增加 P2 输出层
        self.stride = torch.tensor([4, 8, 16, 32])  
        self.no = nc + 5  # 输出维度
        self.nl = 4  # 检测层数
        # ... 其余初始化代码

采用 Focal Loss 替代标准交叉熵损失：

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.8, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred, target):
        bce_loss = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
        pt = torch.exp(-bce_loss)
        loss = self.alpha * (1-pt)**self.gamma * bce_loss
        return loss.mean()

使用 TensorRT 加速的关键步骤：

导出 ONNX 模型：

python export.py --weights yolov5x.pt --include onnx --dynamic

转换 TensorRT 引擎：

import tensorrt as trt

logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open("yolov5x.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
serialized_engine = builder.build_serialized_network(network, config)

性能对比（NVIDIA T4 GPU）：

方案	推理时延 (ms)	mAP@0.5
PyTorch 原生	35.2	72.1%
TensorRT-FP32	18.6	71.9%
TensorRT-FP16	9.4	71.5%

使用可视化工具检查边界框是否准确覆盖目标
统计标注框尺寸分布，过滤异常小（<4px）或异常大（> 图像 50%）的标注
检查类别标签是否正确（特别是负样本）

早停法（Early Stopping）监控验证集 mAP

使用 MixUp 数据增强：

def mixup(imgs, labels, alpha=1.0):
    lam = np.random.beta(alpha, alpha)
    idx = torch.randperm(len(imgs))
    mixed_img = lam * imgs + (1-lam) * imgs[idx]
    mixed_labels = torch.cat([labels, labels[idx]], dim=0)
    return mixed_img, mixed_labels