yolov8数据集缓解样本短缺与类别不均衡【数据增强、过采样、损失加权】

要在数据集不变的情况下缓解 nohelmet 样本短缺问题，可通过 数据增强、过采样、损失加权 三类技术组合优化。以下是具体实现方案与代码示例：

一、核心思路

针对 nohelmet 样本少、模型易“忽视”的问题，从数据多样性（增强）、训练频率（过采样）、损失敏感度（加权）三方面调整，强制模型关注该类别。

二、方法1：数据增强（提升 nohelmet 样本多样性）

通过对 nohelmet 样本施加更强的颜色、几何变换，增加模型对该类的泛化能力。

实现方式：修改 data.yaml 全局增强参数

YOLOv8 支持通过数据集配置文件（data.yaml）调整增强强度，对所有样本生效（nohelmet 因占比低，增强收益更显著）。

# data.yaml 示例（重点修改增强参数）
train: E:/CodeCNN/yolov8-study/VOCdevkit2/dataset2/images/train  # 训练集图像路径
val: E:/CodeCNN/yolov8-study/VOCdevkit2/dataset2/images/val      # 验证集图像路径
names: [helmet, nohelmet, twowheel]  # 类别顺序

# 【增强参数：加大强度，让nohelmet样本更“多变”】
hsv_h: 0.03   # 色相变化范围（默认0.015 → 加大）
hsv_s: 0.8    # 饱和度变化范围（默认0.7 → 加大）
hsv_v: 0.5    # 亮度变化范围（默认0.4 → 加大）
degrees: 20   # 旋转角度（默认0 → 增加）
translate: 0.2 # 平移范围（默认0.1 → 加大）
scale: 0.9    # 缩放范围（默认0.8 → 加大）
shear: 10     # 剪切角度（默认0 → 增加）
flipud: 0.2   # 垂直翻转概率（默认0 → 增加）
fliplr: 0.7   # 水平翻转概率（默认0.5 → 加大）
mosaic: 1.0   # 马赛克增强概率（默认1.0 → 保持，提升小目标鲁棒性）
mixup: 0.3    # 混合增强概率（默认0 → 增加，合成新样本）

进阶：仅对 nohelmet 样本增强（需自定义代码）

若需更精准控制（仅增强 nohelmet），可通过 Albumentations 自定义增强管道，修改 YOLOv8 数据加载逻辑：

from albumentations import Compose, OneOf, HorizontalFlip, RandomBrightnessContrast, ShiftScaleRotate, Blur
from albumentations.core.bbox_utils import BboxParams
from ultralytics.yolo.data import YOLODataset
from ultralytics import YOLO
import yaml

# 1. 读取数据集配置
data_cfg = yaml.safe_load("data.yaml")

# 2. 自定义增强函数（仅对nohelmet样本增强）
def custom_augment(image, bboxes, labels):
    # 筛选nohelmet的标注（假设nohelmet是类别1）
    nohelmet_mask = [label == 1 for label in labels]
    if any(nohelmet_mask):  
        aug = Compose([
            OneOf([  # 随机选一种增强
                HorizontalFlip(p=1.0),          # 水平翻转
                ShiftScaleRotate(shift_limit=0.2, scale_limit=0.3, rotate_limit=30, p=1.0),  # 平移+缩放+旋转
                RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=1.0),  # 亮度+对比度
                Blur(blur_limit=5, p=1.0),      # 模糊
            ], p=0.8),  # 80%概率执行增强
        ], bbox_params=BboxParams(format="yolo", label_fields=["labels"]))  # YOLO格式的bbox
        augmented = aug(image=image, bboxes=bboxes, labels=labels)
        return augmented["image"], augmented["bboxes"], augmented["labels"]
    return image, bboxes, labels

# 3. 继承YOLODataset，替换增强逻辑
class CustomYOLODataset(YOLODataset):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.augment_fn = custom_augment  # 绑定自定义增强

    def __getitem__(self, index):
        img, bboxes, labels = super().__getitem__(index)  # 原始数据加载
        img, bboxes, labels = self.augment_fn(img, bboxes, labels)  # 增强nohelmet
        return img, bboxes, labels

# 4. 构建自定义数据集
train_dataset = CustomYOLODataset(data_cfg, "train", imgsz=800)
val_dataset = YOLODataset(data_cfg, "val", imgsz=800)  # 验证集无需增强

# 5. 训练模型（传入自定义数据集）
model = YOLO("runs/detect/train9/weights/best.pt")  # 加载预训练模型
results = model.train(
    data=data_cfg,
    train_dataset=train_dataset,  # 替换训练集
    val_dataset=val_dataset,      # 验证集保持原逻辑
    epochs=50,
    batch=16,
    imgsz=800,
    # 其他参数（如device、optimizer等）保持默认或按需调整
)

三、方法2：过采样（提升 nohelmet 训练频率）

通过加权采样，让 nohelmet 样本在训练时被更频繁选取，增加模型对该类的学习次数。

实现方式：PyTorch WeightedRandomSampler

import torch
from torch.utils.data import WeightedRandomSampler
from ultralytics.yolo.data import build_dataloader, YOLODataset
from ultralytics import YOLO
import yaml

# 1. 统计训练集各类别实例数
data_cfg = yaml.safe_load("data.yaml")
train_dataset = YOLODataset(data_cfg, "train", imgsz=800)

# 手动统计实例数（helmet=0, nohelmet=1, twowheel=2）
class_counts = [0, 0, 0]
for img_path in train_dataset.im_files:  # 遍历所有训练图像
    label_path = img_path.replace("images", "labels").replace(".jpg", ".txt")
    with open(label_path, "r") as f:
        labels = f.readlines()
    for label in labels:
        cls = int(label.split()[0])
        class_counts[cls] += 1

# 2. 计算采样权重（逆频率加权：样本越少，权重越高）
weights = 1.0 / torch.tensor(class_counts, dtype=torch.float)

# 3. 为每个样本分配权重（按其类别权重）
sample_weights = []
for img_path in train_dataset.im_files:
    label_path = img_path.replace("images", "labels").replace(".jpg", ".txt")
    with open(label_path, "r") as f:
        labels = f.readlines()
    for _ in labels:  # 每个实例对应一个权重
        cls = int(label.split()[0])
        sample_weights.append(weights[cls].item())
sample_weights = torch.tensor(sample_weights, dtype=torch.float)

# 4. 构建加权采样器（replacement=True表示允许重复采样）
sampler = WeightedRandomSampler(sample_weights, len(sample_weights), replacement=True)

# 5. 构建自定义DataLoader
train_loader = build_dataloader(
    data_cfg,
    batch_size=16,
    imgsz=800,
    mode="train",
    sampler=sampler,
    workers=0,  # 根据硬件性能调整（0表示单线程）
)

# 6. 训练模型（传入自定义DataLoader）
model = YOLO("runs/detect/train9/weights/best.pt")
results = model.train(
    data=data_cfg,
    dataloader=train_loader,  # 替换训练DataLoader
    val=True,
    epochs=50,
    imgsz=800,
)

四、方法3：损失加权（让模型“重视” nohelmet 损失）

在计算**分类损失（cls_loss）**时，给 nohelmet 分配更高权重，强制模型关注该类的预测误差。

实现方式：自定义损失函数（需修改YOLOv8训练逻辑）

from ultralytics.yolo.engine.trainer import YOLOTrainer
from ultralytics.yolo.utils.loss import ComputeLoss
import torch
import yaml

class WeightedComputeLoss(ComputeLoss):
    def __init__(self, model, class_weights):
        super().__init__(model)
        self.class_weights = class_weights.to(model.device)  # 类别权重张量（如 [w0, w1, w2]）

    def __call__(self, p, targets):
        loss = super().__call__(p, targets)  # 先计算原始损失
        # 【修改分类损失：乘以类别权重】
        for i in range(self.nl):  # 遍历3个检测头（YOLOv8默认3头）
            # 提取分类预测与目标
            cls_pred = p[i][..., 4:self.nc]  # 分类分支输出 (batch, anchor, grid, grid, nc)
            b, a, gj, gi = self._build_targets(p[i], targets)  # 目标位置索引
            cls_target = torch.zeros_like(cls_pred)
            cls_target[b, a, gj, gi] = 1.0  # 独热编码目标

            # 加权分类损失（BCEWithLogitsLoss）
            cls_loss = self.BCEcls(cls_pred, cls_target) * self.class_weights.view(1, -1, 1, 1)
            loss += cls_loss.sum() * self.cls  # 替换原分类损失
        return loss

class WeightedYOLOTrainer(YOLOTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 统计类别权重（同过采样逻辑）
        data_cfg = yaml.safe_load(self.args.data)
        train_dataset = self.train_loader.dataset
        class_counts = [0, 0, 0]  # 手动统计或调用train_dataset.get_class_counts()
        for img_path in train_dataset.im_files:
            label_path = img_path.replace("images", "labels").replace(".jpg", ".txt")
            with open(label_path, "r") as f:
                labels = f.readlines()
            for label in labels:
                cls = int(label.split()[0])
                class_counts[cls] += 1
        self.class_weights = 1.0 / torch.tensor(class_counts, dtype=torch.float)
        
        # 替换损失函数为加权版
        self.loss = WeightedComputeLoss(self.model, self.class_weights)

# 训练时使用自定义训练器
model = YOLO("runs/detect/train9/weights/best.pt")
trainer = WeightedYOLOTrainer(
    model=model.model,
    data=model.args.data,
    epochs=model.args.epochs,
    batch=model.args.batch,
    imgsz=model.args.imgsz,
    # 其他参数（如device、optimizer等）保持默认
)
trainer.train()

五、组合方案（推荐）

为最大化效果，建议数据增强（全局+自定义）+ 损失加权 组合使用：

1. 在 data.yaml 中加大全局增强强度（覆盖所有样本，nohelmet 受益更明显）。
2. 自定义损失函数，给 nohelmet 分配 class_weights = 1 / nohelmet实例数（如 helmet 实例12000、nohelmet 6000、twowheel 13000，则权重为 [1/12000, 1/6000, 1/13000]，nohelmet 权重是 helmet 的2倍）。

关键参数解释

通过上述方法，可在不新增原始数据的情况下，有效缓解 nohelmet 样本短缺导致的模型偏向问题，提升该类的检测召回率与mAP。