AITOD数据集入门指南：从数据加载到模型训练的全流程解析

1次阅读

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

AITOD 是一个专注于航空图像目标检测的数据集，包含多种航空场景下的目标标注。它的主要特点包括高分辨率图像、小目标密集分布以及复杂的背景干扰。相比通用目标检测数据集（如 COCO），AITOD 更适用于无人机航拍、遥感监测等垂直领域。

数据规模：约 50,000 张图像，覆盖 10 个常见航空目标类别
标注格式：采用 COCO 标准的 JSON 标注文件
独特优势：包含大量小目标样本（平均每图 15+ 个目标），并提供云层遮挡等真实场景标注

安装依赖：
```
pip install pycocotools torchvision
```

基础加载代码（使用 PyTorch 和 COCO API）：

from pycocotools.coco import COCO
import torchvision.transforms as T

# 初始化 COCO API
dataDir = './AITOD'
annFile = f'{dataDir}/annotations/instances_train.json'
coco = COCO(annFile)

# 定义预处理流程
transform = T.Compose([T.ToTensor(),
    T.Resize((800, 800)),
    T.Normalize(mean=[0.485, 0.456, 0.406], 
               std=[0.229, 0.224, 0.225])
])

自定义 Dataset 类：

class AITODDataset(torch.utils.data.Dataset):
    def __init__(self, coco, transform=None):
        self.coco = coco
        self.img_ids = list(coco.imgs.keys())
        self.transform = transform

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]
        img_info = self.coco.loadImgs(img_id)[0]
        img_path = f'{dataDir}/images/{img_info["file_name"]}'
        img = Image.open(img_path).convert('RGB')

        ann_ids = self.coco.getAnnIds(imgIds=img_id)
        anns = self.coco.loadAnns(ann_ids)

        boxes = [ann['bbox'] for ann in anns]  # [x,y,w,h]格式
        labels = [ann['category_id'] for ann in anns]

        if self.transform:
            img = self.transform(img)

        return img, {'boxes': boxes, 'labels': labels}

基础训练配置（使用 Faster R-CNN 示例）：

import torchvision
from torch.utils.data import DataLoader

# 加载预训练模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 10  # AITOD 类别数
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# 准备数据
dataset = AITODDataset(coco, transform=transform)
loader = DataLoader(dataset, batch_size=4, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))

# 训练循环
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9)
for epoch in range(10):
    for images, targets in loader:
        optimizer.zero_grad()
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        losses.backward()
        optimizer.step()

关键训练技巧：
使用学习率预热（Learning Rate Warmup）应对初期不稳定梯度
对小目标检测建议使用 FPN（特征金字塔）结构
数据增强推荐：随机裁剪、色彩抖动（避免几何变形破坏航空图像比例）

问题 1：显存不足
解决方案：减小 batch size（可小至 2）、使用梯度累积
示例代码：每 4 个 mini-batch 执行一次参数更新
问题 2：小目标漏检
解决方案：调整 anchor 大小（建议设置为 16-64 像素范围）
修改 FPN 的 P2 层特征输出
问题 3：标注格式冲突
注意点：AITOD 使用 COCO 格式的[x,y,w,h]，而 PyTorch 需要[x1,y1,x2,y2]

转换代码：

boxes = torch.as_tensor([[x, y, x+w, y+h] for (x,y,w,h) in boxes
], dtype=torch.float32)

数据加载优化：
使用 torch.utils.data.DataLoader 的num_workers=4参数
对图像预处理采用 torchvision.transforms.functional 函数式接口

混合精度训练：

from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()

with autocast():
    loss_dict = model(images, targets)
scaler.scale(losses).backward()
scaler.step(optimizer)
scaler.update()