LangFlow Dynatrace AI驱动运维洞察
2025/12/23 2:21:54
目录
0. 前言——为什么需要自定义可视化？
1. 核心功能点
2. 技术细节：如何实现“红底白字”？
3. 完整代码实现
4. 使用说明
5. 总结
0. 前言——为什么需要自定义可视化？
在使用 Ultralytics YOLO（v8/v9/v10/v11）进行目标检测时，我们通常直接调用 model.predict(save=True)。虽然官方自带的可视化很方便，但它会根据不同的类别分配不同的颜色。
在某些学术论文、工业报告或特定演示场景下，我们往往需要更统一、更专业的视觉风格——例如：亮红色的检测框、红色的标签背景，搭配纯白色的文字。
今天分享一个实用的 Python 脚本，直接绕过官方 UI，利用OpenCV重新绘制：红色边框、红色标签底、纯白文字。让你的检测结果图瞬间达到出版级水准！
1. 核心功能点
极简视觉风格：统一采用“红框 + 红底 + 白字”，重点突出，风格专业。
双重输出：既生成自定义样式的图片，也支持同步保存标准 YOLO 格式的 .txt 标签。
坐标精准处理：自动处理标签背景框越界问题（当物体在图片顶部时，标签会自动调整位置）。
参数化配置：支持通过命令行调整置信度、框粗细、字体大小等。
2. 技术细节：如何实现“红底白字”？
官方的 plot() 函数高度封装，难以精细修改颜色逻辑。本脚本通过以下逻辑实现：
提取结果：利用 result.boxes 获取坐标、类别和置信度。
OpenCV 绘制：
cv2.rectangle：画出红色（BGR: 0, 0, 255）的边框。
cv2.getTextSize：计算文字宽度，动态生成背景红色矩形。
cv2.putText：在红色背景上叠写白色文字。
3. 完整代码实现
您可以直接将以下代码保存为 custom_predict.py 并运行。
import argparse import os import warnings import cv2 from ultralytics import YOLO warnings.filterwarnings("ignore") def save_txt_yolo_format(result, img_w, img_h, txt_path, save_conf=True): """ 保存 YOLO 格式标签: class x_center y_center w h [conf] 归一化到 0~1 """ lines = [] if result.boxes is None or len(result.boxes) == 0: open(txt_path, "w").close() return boxes_xyxy = result.boxes.xyxy.cpu().numpy() clss = result.boxes.cls.cpu().numpy().astype(int) confs = result.boxes.conf.cpu().numpy() if result.boxes.conf is not None else None for i, (x1, y1, x2, y2) in enumerate(boxes_xyxy): x1, y1, x2, y2 = float(x1), float(y1), float(x2), float(y2) xc = (x1 + x2) / 2.0 / img_w yc = (y1 + y2) / 2.0 / img_h bw = (x2 - x1) / img_w bh = (y2 - y1) / img_h if save_conf and confs is not None: lines.append(f"{clss[i]} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f} {confs[i]:.6f}") else: lines.append(f"{clss[i]} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}") with open(txt_path, "w") as f: f.write("\n".join(lines)) def draw_red_box_red_bg_white_text(results, out_dir, save_txt=False, save_conf=True, line_thickness=2, font_scale=0.6): """ 框=红、标签底=红、文字=白色（类别+置信度） """ os.makedirs(out_dir, exist_ok=True) labels_dir = os.path.join(out_dir, "labels") if save_txt: os.makedirs(labels_dir, exist_ok=True) red = (0, 0, 255) # BGR 红色 white = (255, 255, 255) # BGR 白色 for r in results: img_path = r.path img = cv2.imread(img_path) if img is None: print(f"⚠️ 读图失败: {img_path}") continue h, w = img.shape[:2] names = r.names # dict: class_id -> class_name # 保存 txt if save_txt: base = os.path.splitext(os.path.basename(img_path))[0] txt_path = os.path.join(labels_dir, base + ".txt") save_txt_yolo_format(r, w, h, txt_path, save_conf=save_conf) # 画框 + 标签 if r.boxes is not None and len(r.boxes) > 0: boxes = r.boxes.xyxy.cpu().numpy() confs = r.boxes.conf.cpu().numpy() if r.boxes.conf is not None else None clss = r.boxes.cls.cpu().numpy().astype(int) if r.boxes.cls is not None else None for i, (x1, y1, x2, y2) in enumerate(boxes): x1, y1, x2, y2 = map(int, [x1, y1, x2, y2]) # 红框 cv2.rectangle(img, (x1, y1), (x2, y2), red, line_thickness) # 标签文本：class + conf label = "" if clss is not None: label += names.get(int(clss[i]), str(int(clss[i]))) if save_conf and confs is not None: label += f" {confs[i]:.2f}" if label.strip(): (tw, th), baseline = cv2.getTextSize( label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 2 ) # 标签放在框上方，避免越界；放不下就顶到0 y_top = max(0, y1 - th - baseline - 6) # 标签背景：红色 cv2.rectangle( img, (x1, y_top), (x1 + tw + 6, y1), red, -1 ) # 标签文字：白色 cv2.putText( img, label, (x1 + 3, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, font_scale, white, 2, cv2.LINE_AA ) save_path = os.path.join(out_dir, os.path.basename(img_path)) cv2.imwrite(save_path, img) def main(opt): model = YOLO(opt.weights) print(f"🔍 开始预测：{opt.source}") print(f"✅ 权重：{opt.weights}") # 预测：关闭官方保存，自己画（才能控颜色） results = model.predict( source=opt.source, imgsz=opt.img_size, device=opt.device, conf=opt.conf, save=False, verbose=False ) out_dir = os.path.join(opt.project, opt.name) draw_red_box_red_bg_white_text( results, out_dir=out_dir, save_txt=opt.save_txt, save_conf=True, line_thickness=opt.thickness, font_scale=opt.font_scale ) print(f"✅ 完成（红框+红底+白字），输出目录：{out_dir}") if opt.save_txt: print(f"✅ txt 已保存到：{os.path.join(out_dir, 'labels')}") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--weights", type=str, default="/root/autodl-tmp/ultralytics-main/runs/detect/train3/weights/best.pt", help="model weight path" ) parser.add_argument( "--source", type=str, default="/root/autodl-tmp/ultralytics-main/images", help="image dir or image file" ) parser.add_argument("--img_size", type=int, default=640, help="image size") parser.add_argument("--device", type=str, default="0", help="cuda device") parser.add_argument("--conf", type=float, default=0.25, help="confidence threshold") parser.add_argument("--project", type=str, default="runs/predict", help="save root dir") parser.add_argument("--name", type=str, default="train1_pred", help="save sub dir name") parser.add_argument("--save_txt", action="store_true", help="save yolo txt labels") parser.add_argument("--thickness", type=int, default=2, help="box line thickness") parser.add_argument("--font_scale", type=float, default=0.6, help="label font scale") opt = parser.parse_args() main(opt)4. 使用说明
1. 基础预测：
codeBash
downloadcontent_copy
expand_less
python custom_predict.py --weights ./runs/train/best.pt --source ./data/test_images2. 调整线条粗细与字体（适合高分辨率图）：
codeBash
downloadcontent_copy
expand_less
python custom_predict.py --weights best.pt --source ./images --thickness 4 --font_scale 0.83. 同时生成标签文件（用于辅助标注）：
codeBash
downloadcontent_copy
expand_less
python custom_predict.py --weights best.pt --source ./images --save_txt5. 总结
通过简单的 OpenCV 二次开发，我们可以轻松突破 YOLO 官方库的可视化限制。这种红底白字的风格不仅统一了视觉语言，更在复杂背景下提供了极佳的辨识度。
如果你觉得这个脚本有用，欢迎点赞、收藏，并关注我的 CSDN 账号，获取更多深度学习实用工具！