如何自定义YOLOv5推理阶段的标签字体、Bounding Box颜色及粗细

一、自定义Bounding Box颜色 #

YOLOv5的results.render()方法支持传入colors参数，为每个类别指定专属的RGB颜色值。具体操作如下：

1. 获取模型的类别列表
2. 为每个类别定义对应的RGB颜色（颜色数量需与类别数一致）
3. 将颜色列表传入render()方法

修改后的示例代码：

import torch
import cv2
import numpy as np

model = torch.hub.load('ultralytics/yolov5', 'custom', path=r'C:\Users\rohin\Desktop\NewData\yolo training\last.pt', force_reload=True)
class_names = model.names  # 获取所有类别名称
# 为每个类别定义RGB颜色，示例中3个类别对应3种颜色，按需调整
custom_colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]  

cap = cv2.VideoCapture(0)
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 执行检测
    results = model(frame)
    # 传入自定义颜色渲染结果
    rendered_frame = np.squeeze(results.render(colors=custom_colors))

    cv2.imshow('YOLO', rendered_frame)
    if cv2.waitKey(10) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

二、自定义标签字体 #

默认的results.render()方法封装了绘制逻辑，无法直接修改字体。需要手动解析检测结果，使用OpenCV的绘图API自定义绘制框和标签，支持修改字体类型、大小、厚度等参数：

操作步骤： #

1. 从results.pred[0]中提取检测框坐标、置信度、类别ID
2. 选择OpenCV内置字体（如cv2.FONT_HERSHEY_COMPLEX、cv2.FONT_HERSHEY_TRIPLEX等），或通过PIL加载自定义TTF字体
3. 遍历检测结果，分别绘制Bounding Box和标签文本

完整示例代码： #

import torch
import cv2
import numpy as np

model = torch.hub.load('ultralytics/yolov5', 'custom', path=r'C:\Users\rohin\Desktop\NewData\yolo training\last.pt', force_reload=True)
class_names = model.names
# 自定义类别颜色，按需调整
custom_colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]
# 选择自定义字体：这里用OpenCV内置的复杂字体，也可替换为其他内置字体
custom_font = cv2.FONT_HERSHEY_COMPLEX
font_scale = 0.8  # 字体大小
font_thickness = 2  # 字体厚度

cap = cv2.VideoCapture(0)
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 执行检测
    results = model(frame)
    detections = results.pred[0]  # 获取检测结果，每个元素格式：(x1, y1, x2, y2, conf, cls)

    # 遍历所有检测结果，手动绘制
    for det in detections:
        x1, y1, x2, y2, conf, cls = det
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        cls_idx = int(cls)
        # 绘制Bounding Box
        cv2.rectangle(frame, (x1, y1), (x2, y2), custom_colors[cls_idx], 2)
        # 生成标签文本
        label = f'{class_names[cls_idx]} {conf:.2f}'
        # 获取文本尺寸，用于调整标签位置
        (text_w, text_h), _ = cv2.getTextSize(label, custom_font, font_scale, font_thickness)
        # 绘制标签背景（可选，提升可读性）
        cv2.rectangle(frame, (x1, y1 - text_h - 10), (x1 + text_w, y1), custom_colors[cls_idx], -1)
        # 绘制标签文本
        cv2.putText(frame, label, (x1, y1 - 5), custom_font, font_scale, (255, 255, 255), font_thickness)

    cv2.imshow('YOLO', frame)
    if cv2.waitKey(10) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

加载自定义TTF字体（可选） #

如果需要使用系统中的TTF字体，可通过PIL库配合OpenCV实现：

from PIL import ImageFont, ImageDraw, Image

# 加载自定义TTF字体，替换为你的字体路径
font_path = 'C:/Windows/Fonts/msyh.ttc'
custom_font = ImageFont.truetype(font_path, 20)  # 20为字体大小

# 在绘制标签时替换为以下逻辑：
frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
draw = ImageDraw.Draw(frame_pil)
draw.text((x1, y1 - text_h - 10), label, font=custom_font, fill=(255,255,255))
frame = cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR)

内容的提问来源于stack exchange，提问作者rohin