📚 课程目录第1章 YOLOv5环境搭建与COCO数据集准备

1. YOLOv5环境搭建与COCO数据集准备

任务目标

本案例完成YOLOv5训练环境搭建，下载COCO数据集的少量子集（person类别），转换为YOLO格式并验证数据加载流程。最终通过运行数据检查脚本，确认图像与标签文件正确配对、类别ID无误。

环境依赖

需要Python 3.8+、PyTorch 1.8+、Ultralytics YOLOv5仓库。建议使用GPU（CUDA 11.x），CPU也可运行但训练较慢。

# 创建虚拟环境并安装依赖
python -m venv yolov5_env
source yolov5_env/bin/activate  # Windows: yolov5_env\Scripts\activate
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

COCO子集下载与转换

使用pycocotools从COCO验证集中筛选仅含person类别的图像（约100张），并转换为YOLO格式的标签文件。以下脚本完成下载、筛选、格式转换与数据集划分：

import os
import json
import shutil
import urllib.request
import zipfile
from pathlib import Path
from collections import defaultdict

# 配置路径
BASE_DIR = Path("./coco_person_subset")
IMG_DIR = BASE_DIR / "images"
LABEL_DIR = BASE_DIR / "labels"
ANNOTATION_URL = "http://images.cocodataset.org/annotations/annotations_trainval2017.zip"
IMAGE_BASE_URL = "http://images.cocodataset.org/val2017/"

# 创建目录
IMG_DIR.mkdir(parents=True, exist_ok=True)
LABEL_DIR.mkdir(parents=True, exist_ok=True)

# 1. 下载COCO标注文件（若未存在）
ann_zip = BASE_DIR / "annotations_trainval2017.zip"
if not ann_zip.exists():
    print("下载COCO标注文件...")
    urllib.request.urlretrieve(ANNOTATION_URL, ann_zip)
    with zipfile.ZipFile(ann_zip, 'r') as z:
        z.extractall(BASE_DIR)

# 2. 加载instances_val2017.json
with open(BASE_DIR / "annotations" / "instances_val2017.json", 'r') as f:
    coco_data = json.load(f)

# 建立类别名到ID的映射
cat_id_to_name = {cat['id']: cat['name'] for cat in coco_data['categories']}
cat_name_to_id = {cat['name']: cat['id'] for cat in coco_data['categories']}
person_cat_id = cat_name_to_id['person']

# 3. 筛选包含person的图像
img_id_to_info = {img['id']: img for img in coco_data['images']}
img_ids_with_person = set()
annotations_by_img = defaultdict(list)

for ann in coco_data['annotations']:
    if ann['category_id'] == person_cat_id:
        img_ids_with_person.add(ann['image_id'])
        annotations_by_img[ann['image_id']].append(ann)

# 取前100张包含person的图像
selected_img_ids = list(img_ids_with_person)[:100]
print(f"筛选到 {len(selected_img_ids)} 张包含person的图像")

# 4. 下载图像并生成YOLO格式标签
for idx, img_id in enumerate(selected_img_ids):
    img_info = img_id_to_info[img_id]
    file_name = img_info['file_name']
    width, height = img_info['width'], img_info['height']

    # 下载图像
    img_path = IMG_DIR / file_name
    if not img_path.exists():
        urllib.request.urlretrieve(IMAGE_BASE_URL + file_name, img_path)

    # 生成YOLO标签文件
    label_path = LABEL_DIR / (Path(file_name).stem + ".txt")
    with open(label_path, 'w') as f:
        for ann in annotations_by_img[img_id]:
            # COCO bbox: [x, y, w, h] → YOLO: [x_center, y_center, w, h] 归一化
            x, y, w, h = ann['bbox']
            x_center = (x + w / 2) / width
            y_center = (y + h / 2) / height
            w_norm = w / width
            h_norm = h / height
            # person类别ID在YOLO中为0（COCO person的cat_id=1，但YOLO从0开始）
            f.write(f"0 {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}\n")

    if (idx + 1) % 20 == 0:
        print(f"已处理 {idx + 1}/{len(selected_img_ids)} 张图像")

# 5. 生成数据集配置文件
data_yaml = f"""
path: {BASE_DIR.resolve()}
train: images
val: images
nc: 1
names: ['person']
"""
with open(BASE_DIR / "dataset.yaml", 'w') as f:
    f.write(data_yaml)

print("数据集准备完成！")
print(f"图像数量: {len(list(IMG_DIR.glob('*.jpg')))}")
print(f"标签数量: {len(list(LABEL_DIR.glob('*.txt')))}")

数据加载验证

运行以下脚本验证YOLOv5能否正确加载数据集并读取标签：

import torch
from pathlib import Path

# 加载YOLOv5数据加载器
from utils.datasets import LoadImagesAndLabels

data_yaml = Path("./coco_person_subset/dataset.yaml")
dataset = LoadImagesAndLabels(str(data_yaml), img_size=640, batch_size=8, augment=False)

# 验证第一个batch
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True)
batch = next(iter(dataloader))
imgs, labels, paths, shapes = batch

print(f"Batch图像张量形状: {imgs.shape}")  # 预期: [8, 3, 640, 640]
print(f"Batch标签形状: {labels.shape}")    # 预期: [N, 6] 每行: [image_idx, cls, x, y, w, h]
print(f"样本图像路径: {paths[:2]}")

# 检查标签是否在有效范围内
cls = labels[:, 1]
assert cls.max() == 0, "类别ID应为0（person）"
assert (labels[:, 2:] >= 0).all() and (labels[:, 2:] <= 1).all(), "归一化坐标应在[0,1]内"
print("数据加载验证通过！")

验证方法与判定标准

数据集准备验证：运行第一个脚本后，终端输出图像数量与标签数量均为100，且dataset.yaml文件生成成功。
数据加载验证：运行第二个脚本，若输出Batch图像张量形状: [8, 3, 640, 640]且无断言错误，说明数据加载流程正常。
标签正确性抽查：随机打开coco_person_subset/labels/下任一txt文件，确认每行格式为0 x_center y_center width height，数值均在0~1之间。