bit-vehicle数据集实战指南：从数据预处理到模型训练的全流程解析

1次阅读

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

bit-vehicle 数据集是由北京理工大学发布的车辆图像数据集，包含 9851 张高分辨率车辆图片，涵盖 6 种车型（SUV、轿车、小型客车等），每张图片都有精确的边界框标注。该数据集特别适合车辆识别、智能交通监控等场景，但因采集自真实道路环境，存在光照变化、遮挡等挑战。

标注格式特殊 ：原始标注为.mat 文件，需转换为常用格式（如 COCO 或 Pascal VOC）
图像尺寸不一 ：分辨率从 640×480 到 1920×1080 不等，需统一处理
类别不平衡 ：某些车型样本量不足（如卡车仅占 5%）
标注错误 ：少量边界框存在漏标或错标情况

import scipy.io
import json
import os

# 转换.mat 标注为 COCO 格式
def convert_to_coco(mat_path, output_dir):
    data = scipy.io.loadmat(mat_path)
    annotations = []

    for idx, (img_name, boxes) in enumerate(zip(data['vehicles'][0], data['annotation'][0])):
        image_id = idx + 1
        for box in boxes:
            x1, y1, x2, y2 = box[0][0], box[1][0], box[2][0], box[3][0]
            annotations.append({
                'image_id': image_id,
                'bbox': [x1, y1, x2-x1, y2-y1],  # COCO 格式要求 [x,y,width,height]
                'category_id': int(box[4][0]) - 1  # 原标签从 1 开始
            })

    with open(os.path.join(output_dir, 'annotations.json'), 'w') as f:
        json.dump({'annotations': annotations}, f)

from PIL import Image
import numpy as np

def resize_with_padding(img, target_size=(640, 640)):
    # 保持宽高比的缩放 + 填充
    original_ratio = img.width / img.height
    target_ratio = target_size[0] / target_size[1]

    if original_ratio > target_ratio:
        new_width = target_size[0]
        new_height = int(target_size[0] / original_ratio)
    else:
        new_height = target_size[1]
        new_width = int(target_size[1] * original_ratio)

    img = img.resize((new_width, new_height))
    new_img = Image.new('RGB', target_size, (128, 128, 128))
    new_img.paste(img, ((target_size[0]-new_width)//2, 
                        (target_size[1]-new_height)//2))
    return new_img

import torch
import torchvision
from torch import nn

class VehicleDetector(nn.Module):
    def __init__(self, num_classes=6):
        super().__init__()
        # 使用预训练的 ResNet 作为 backbone
        self.backbone = torchvision.models.resnet18(pretrained=True)

        # 修改最后一层卷积输出通道数
        self.backbone.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)

        # 检测头
        self.head = nn.Sequential(nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, num_classes + 4, kernel_size=3, padding=1)  # 4 对应 bbox 坐标
        )

    def forward(self, x):
        features = self.backbone(x)
        return self.head(features)

推荐使用 Albumentations 库实现针对车辆数据的特殊增强：

import albumentations as A

transform = A.Compose([A.RandomBrightnessContrast(p=0.5),  # 应对光照变化
    A.HorizontalFlip(p=0.5),
    A.RandomRain(p=0.1),  # 模拟雨天场景
    A.RandomShadow(p=0.3),
    A.Cutout(max_h_size=30, max_w_size=30, p=0.2)  # 模拟遮挡
], bbox_params=A.BboxParams(format='coco'))

# 使用加权随机采样
from torch.utils.data import WeightedRandomSampler

class_counts = [1200, 3500, 800, 5000, 200, 150]  # 各类别样本数
weights = 1. / torch.tensor(class_counts, dtype=torch.float)
sampler = WeightedRandomSampler(weights, len(dataset))