联合目标检测与图像分类提升数据不平衡场景下的准确率

在一些数据不平衡的场景下，使用单一的目标检测模型很难达到99%的准确率。为了优化这一问题，适当将其拆解为目标检测模型和图像分类模型的组合，可以更有效地控制最终效果，尤其是在添加焦点损失（focal loss）、调整超参数和数据预处理无效的情况下。以下是具体的实现方式及联合两个模型的推理代码。

整体功能概述

这段代码的主要功能包括：

加载目标检测和分类模型：使用两个 Ultralytics YOLO（YOLOv8/YOLOv11均可）模型进行目标检测和分类。
处理图像：遍历指定输入文件夹中的所有图像，进行目标检测和分类。
绘制检测框和分类标签：在图像上绘制检测到的对象的边界框，并在框上方添加分类名称和置信度。
可选保存裁剪的对象图像：根据设置，裁剪检测到的对象区域并保存为单独的图像文件，文件名包含类别名称、置信度和坐标信息（便于调试）。

实现细节

1. 加载模型

代码加载了两个 YOLO 模型：

目标检测模型：一个单一类别的 YOLO 模型，用于检测主体对象。
图像分类模型：一个多类别的 YOLO 模型，用于对检测到的对象进行分类。

2. 处理图像

脚本处理输入文件夹中的每一张图像，步骤如下：

目标检测：使用目标检测模型检测图像中的对象。
裁剪检测到的对象：根据检测到的边界框坐标，裁剪出感兴趣的区域。
图像分类：对裁剪出的对象区域进行分类。
数据增强或欠采样：根据任务需求，对裁剪出的子图像进行数据增强或欠采样，以平衡数据集。

3. 绘制检测框和标签

对于每一个检测到的对象，脚本会：

在图像上绘制一个边界框。
在边界框上方添加分类名称和置信度标签。

4. 保存裁剪的对象图像

可选地，脚本会保存裁剪出的对象图像，文件名包含以下信息：

分类名称
置信度
边界框坐标

这对于调试和分析特定的检测结果非常有帮助。

推理代码

import os
import cv2
import numpy as np
from pathlib import Path
from ultralytics import YOLO
import randomdef generate_random_color_from_name(name):"""根据类别名生成可重复的颜色。"""random.seed(name)  # 使用类别名作为随机种子return tuple(random.randint(0, 255) for _ in range(3))def generate_class_colors(names):"""为每个类别生成一个固定的颜色。"""class_colors = {}for class_name in names:class_colors[class_name] = generate_random_color_from_name(class_name)return class_colorsdef draw_box_on_image(image, box, color=(0, 255, 0), thickness=2):"""在图像上绘制检测框。"""x1, y1, x2, y2 = map(int, box)cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)def add_classification_to_box(image, box, class_name, confidence, color=(0, 255, 0)):"""在边界框上方添加分类名称和置信度。"""x1, y1, x2, y2 = map(int, box)label = f"{class_name}: {confidence:.2f}"cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2, cv2.LINE_AA)def save_cropped_object(image, box, cls_class_name, confidence, output_folder, image_name):"""将裁剪的对象区域保存为图像到子文件夹中，文件名包含类别名、置信度和坐标。"""x1, y1, x2, y2 = map(int, box)cropped_img = image[y1:y2, x1:x2]# 为当前图像创建一个以图像文件名命名的子文件夹image_subfolder = Path(output_folder) / Path(image_name).stemimage_subfolder.mkdir(parents=True, exist_ok=True)# 为裁剪的对象创建文件名（class_name_confidence_x1_y1_x2_y2.jpg）# 确保置信度格式安全，使用两位小数，并用下划线分隔cropped_img_name = f"{cls_class_name}_{confidence:.2f}_{x1}_{y1}_{x2}_{y2}.jpg"cropped_img_path = image_subfolder / cropped_img_namecv2.imwrite(str(cropped_img_path), cropped_img)print(f"已保存裁剪对象: {cropped_img_path}")def process_image_with_detection_and_classification(model_det, model_cls, img_path, names, class_colors, output_folder, save_cropped=False, detection_size=1280, classification_size=640):"""处理单张图像：执行对象检测，分类每个对象，并返回处理后的图像。:param model_det: 检测模型:param model_cls: 分类模型:param img_path: 图像路径:param names: 类别名称列表:param class_colors: 类别颜色映射字典:param output_folder: 输出文件夹路径:param save_cropped: 是否保存裁剪的对象图像:param detection_size: 检测模型输入图像大小:param classification_size: 分类模型输入图像大小:return: 处理后的图像"""img = cv2.imread(str(img_path))if img is None:print(f"无法读取图像: {img_path}")return None# 创建图像副本用于绘制（不修改原始图像）img_copy = img.copy()# 执行对象检测results_det = model_det.predict(str(img_path), imgsz=detection_size, conf=0.25, iou=0.45)# 处理每个检测结果（每个检测框）for r in results_det:boxes = r.boxes.xyxy.cpu().numpy()  # xyxy 格式classes = r.boxes.cls.cpu().numpy()confidences = r.boxes.conf.cpu().numpy()for box, cls_id, confidence in zip(boxes, classes, confidences):# 检测模型的类别名det_class_name = names[int(cls_id)]# 使用检测到的类别名对应的颜色（该颜色是全局唯一的）color = class_colors.get(det_class_name, (255, 255, 255))# 裁剪对象区域x1, y1, x2, y2 = map(int, box)object_region = img[y1:y2, x1:x2]# 将对象区域调整为分类模型的输入大小object_region = cv2.resize(object_region, (classification_size, classification_size))# 执行分类results_cls = model_cls.predict(object_region, imgsz=classification_size)for result in results_cls:try:# 获取Top1预测结果classification_confidence = result.probs.cpu().numpy().top1conftop1_index = result.probs.top1cls_class_name = names[top1_index]# 根据分类结果的类别名设置颜色final_color = class_colors.get(cls_class_name, color)add_classification_to_box(img_copy, box, cls_class_name, classification_confidence, color=final_color)# 如果启用了保存裁剪对象，则保存if save_cropped:save_cropped_object(img, box, cls_class_name, classification_confidence, output_folder, img_path.name)except Exception as e:print(f"分类时出错: {e}")# 在图像副本上绘制检测框draw_box_on_image(img_copy, box, color=color)return img_copydef process_images(model_det, model_cls, input_folder, output_folder, names, class_colors, save_cropped=False, detection_size=1280, classification_size=640):"""处理输入文件夹中的图像，执行对象检测和分类，并保存处理后的图像。:param model_det: 检测模型:param model_cls: 分类模型:param input_folder: 输入文件夹路径:param output_folder: 输出文件夹路径:param names: 类别名称列表:param class_colors: 类别颜色映射字典:param save_cropped: 是否保存裁剪的对象图像:param detection_size: 检测模型输入图像大小:param classification_size: 分类模型输入图像大小"""Path(output_folder).mkdir(parents=True, exist_ok=True)image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.webp']for ext in image_extensions:for img_path in Path(input_folder).glob(ext):print(f"正在处理: {img_path}")processed_img = process_image_with_detection_and_classification(model_det, model_cls, img_path, names, class_colors, output_folder, save_cropped, detection_size, classification_size)if processed_img is not None:output_image_path = Path(output_folder) / f"{img_path.stem}_with_boxes_and_classification.jpg"cv2.imwrite(str(output_image_path), processed_img)print(f"已保存处理后的图像: {output_image_path}")else:print(f"跳过图像: {img_path} (无法处理)")if __name__ == '__main__':# 设置是否保存裁剪的对象图像（默认不保存）SAVE_CROPPED = True  # 设置为 True 以启用保存裁剪对象# 加载检测和分类模型model_det = YOLO('runs/device_train/exp9/weights/best.pt')model_cls = YOLO('runs/cls_99.4%_exp14/weights/best.pt')# 设置输入和输出文件夹路径input_folder = 'test1'output_folder = 'infer-1216'# 获取类别名（用于生成一致的类别颜色映射）# 这里使用一张全白的图像来获取类别名black_image = 255 * np.ones((224, 224, 3), dtype=np.uint8)results = model_cls.predict(source=black_image)name_dict = results[0].namesnames = list(name_dict.values())# 只在这里生成一次类别颜色映射class_colors = generate_class_colors(names)# 开始处理图像process_images(model_det, model_cls, input_folder, output_folder,names, class_colors,save_cropped=SAVE_CROPPED,detection_size=1280,classification_size=224)

执行完后的结果
在这里插入图片描述

下面贴一下目标检测和图像分类的ultralytics的训练代码

目标检测训练代码

注意把single_cls=False改成True，变成单类训练

# nohup python -m torch.distributed.launch --nproc_per_node=4 --master_port=25643 det_train.py > output-lane-1212.txt 2>&1 &
# nohup python -m torch.distributed.launch --nproc_per_node=5 --master_port=25698 det_train.py > output-lane-1212.txt 2>&1 &
from ultralytics import YOLOif __name__ == '__main__':# 加载模型model = YOLO("checkpoints/yolo11l.pt")  # 使用预训练权重训练# 训练参数 ----------------------------------------------------------------------------------------------model.train(data='/home/lizhijun/01.det/ultralytics-8.3.23/datasets/device_1212_yolo_without_vdd/config.yaml',epochs=150,  # (int) 训练的周期数patience=50,  # (int) 等待无明显改善以进行早期停止的周期数batch=16,  # (int) 每批次的图像数量（-1 为自动批处理）imgsz=1280,  # (int) 输入图像的大小，整数或w，hsave=True,  # (bool) 保存训练检查点和预测结果save_period=-1,  # (int) 每x周期保存检查点（如果小于1则禁用）cache=False,  # (bool) True/ram、磁盘或False。使用缓存加载数据device='1,2,3,5',  # (int | str | list, optional) 运行的设备，例如 cuda device=0 或 device=0,1,2,3 或 device=cpuworkers=8,  # (int) 数据加载的工作线程数（每个DDP进程）project='runs/device_train',  # (str, optional) 项目名称name='exp',  # (str, optional) 实验名称，结果保存在'project/name'目录下exist_ok=False,  # (bool) 是否覆盖现有实验pretrained=True,  # (bool | str) 是否使用预训练模型（bool），或从中加载权重的模型（str）optimizer='auto',  # (str) 要使用的优化器，选择=[SGD，Adam，Adamax，AdamW，NAdam，RAdam，RMSProp，auto]verbose=True,  # (bool) 是否打印详细输出seed=0,  # (int) 用于可重复性的随机种子deterministic=True,  # (bool) 是否启用确定性模式single_cls=False,  # (bool) 将多类数据训练为单类rect=False,  # (bool) 如果mode='train'，则进行矩形训练，如果mode='val'，则进行矩形验证cos_lr=True,  # (bool) 使用余弦学习率调度器close_mosaic=10,  # (int) 在最后几个周期禁用马赛克增强resume=False,  # (bool) 从上一个检查点恢复训练amp=True,  # (bool) 自动混合精度（AMP）训练，选择=[True, False]，True运行AMP检查fraction=1.0,  # (float) 要训练的数据集分数（默认为1.0，训练集中的所有图像）profile=False,  # (bool) 在训练期间为记录器启用ONNX和TensorRT速度freeze= None,  # (int | list, 可选) 在训练期间冻结前 n 层，或冻结层索引列表。# 超参数 ----------------------------------------------------------------------------------------------lr0=0.01,  # (float) 初始学习率（例如，SGD=1E-2，Adam=1E-3）lrf=0.01,  # (float) 最终学习率（lr0 * lrf）momentum=0.937,  # (float) SGD动量/Adam beta1weight_decay=0.0005,  # (float) 优化器权重衰减 5e-4warmup_epochs=3.0,  # (float) 预热周期（分数可用）warmup_momentum=0.8,  # (float) 预热初始动量warmup_bias_lr=0.1,  # (float) 预热初始偏置学习率box=6,  # (float) 盒损失增益cls=1.5,  # (float) 类别损失增益（与像素比例）dfl=1.5,  # (float) dfl损失增益pose=12.0,  # (float) 姿势损失增益kobj=1.0,  # (float) 关键点对象损失增益label_smoothing=0.05,  # (float) 标签平滑（分数）nbs=64,  # (int) 名义批量大小hsv_h=0.015,  # (float) 图像HSV-Hue增强（分数）hsv_s=0.7,  # (float) 图像HSV-Saturation增强（分数）hsv_v=0.4,  # (float) 图像HSV-Value增强（分数）degrees=90.0,  # (float) 图像旋转（+/- deg）translate=0.5,  # (float) 图像平移（+/- 分数）scale=0.5,  # (float) 图像缩放（+/- 增益）shear=0.4,  # (float) 图像剪切（+/- deg）perspective=0.0,  # (float) 图像透视（+/- 分数），范围为0-0.001flipud=0.5,  # (float) 图像上下翻转（概率）fliplr=0.5,  # (float) 图像左右翻转（概率）mosaic=1.0,  # (float) 图像马赛克（概率）mixup=0.0,  # (float) 图像混合（概率）copy_paste=0.0,  # (float) 分割复制-粘贴（概率）)

图像分类训练代码

from ultralytics import YOLOmodel = YOLO("checkpoints/yolo11l-cls.pt")
model.train(data='/home/lizhijun/01.det/ultralytics-8.3.23/datasets/device_cls_merge_manual_with_21w_1218_train_val_224_truncate_grid_110%', project='runs/cls_train',  # (str, optional) 项目名称name='exp',  # (str, optional) 实验名称，结果保存在'project/name'目录下epochs=20, batch=1024,device='1,2,3,5',erasing=0.0,crop_fraction=1.0,augment=False,auto_augment=False,hsv_h=0.015,  # (float) 图像HSV-Hue增强（分数）hsv_s=0.7,  # (float) 图像HSV-Saturation增强（分数）hsv_v=0.4,  # (float) 图像HSV-Value增强（分数）degrees=0.0,  # (float) 图像旋转（+/- deg）translate=0.0,  # (float) 图像平移（+/- 分数）scale=0.0,  # (float) 图像缩放（+/- 增益）shear=0.0,  # (float) 图像剪切（+/- deg）perspective=0.0,  # (float) 图像透视（+/- 分数），范围为0-0.001flipud=0.5,  # (float) 图像上下翻转（概率）fliplr=0.5,  # (float) 图像左右翻转（概率）mosaic=1.0,  # (float) 图像马赛克（概率）mixup=0.0)  # (float) 图像混合（概率）)