联合目标检测与图像分类提升数据不平衡场景下的准确率
在一些数据不平衡的场景下，使用单一的目标检测模型很难达到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 random
def 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_colors
def 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).stem
image_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_name
cv2.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().top1conf
top1_index = result.probs.top1
cls_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_copy
def 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].names
names = 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 YOLO
if __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，h
save=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=cpu
workers=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 beta1
weight_decay=0.0005, # (float) 优化器权重衰减 5e-4
warmup_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.001
flipud=0.5, # (float) 图像上下翻转（概率）
fliplr=0.5, # (float) 图像左右翻转（概率）
mosaic=1.0, # (float) 图像马赛克（概率）
mixup=0.0, # (float) 图像混合（概率）
copy_paste=0.0, # (float) 分割复制-粘贴（概率）
)
图像分类训练代码
from ultralytics import YOLO
model = 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.001
flipud=0.5, # (float) 图像上下翻转（概率）
fliplr=0.5, # (float) 图像左右翻转（概率）
mosaic=1.0, # (float) 图像马赛克（概率）
mixup=0.0) # (float) 图像混合（概率）)
