当前位置: 首页 > news >正文

【程序干货】YOLO 预测检测结果不顺心?手把手教你自定义“红框白字”专业视觉效果(附完整代码)

目录

0. 前言——为什么需要自定义可视化?

1. 核心功能点

2. 技术细节:如何实现“红底白字”?

3. 完整代码实现

4. 使用说明

5. 总结


0. 前言——为什么需要自定义可视化?

在使用 Ultralytics YOLO(v8/v9/v10/v11)进行目标检测时,我们通常直接调用 model.predict(save=True)。虽然官方自带的可视化很方便,但它会根据不同的类别分配不同的颜色。

在某些学术论文、工业报告特定演示场景下,我们往往需要更统一、更专业的视觉风格——例如:亮红色的检测框、红色的标签背景,搭配纯白色的文字

今天分享一个实用的 Python 脚本,直接绕过官方 UI,利用OpenCV重新绘制:红色边框、红色标签底、纯白文字。让你的检测结果图瞬间达到出版级水准!

1. 核心功能点

  1. 极简视觉风格:统一采用“红框 + 红底 + 白字”,重点突出,风格专业。

  2. 双重输出:既生成自定义样式的图片,也支持同步保存标准 YOLO 格式的 .txt 标签。

  3. 坐标精准处理:自动处理标签背景框越界问题(当物体在图片顶部时,标签会自动调整位置)。

  4. 参数化配置:支持通过命令行调整置信度、框粗细、字体大小等。


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_images

2. 调整线条粗细与字体(适合高分辨率图):

codeBash

downloadcontent_copy

expand_less

python custom_predict.py --weights best.pt --source ./images --thickness 4 --font_scale 0.8

3. 同时生成标签文件(用于辅助标注):

codeBash

downloadcontent_copy

expand_less

python custom_predict.py --weights best.pt --source ./images --save_txt

5. 总结

通过简单的 OpenCV 二次开发,我们可以轻松突破 YOLO 官方库的可视化限制。这种红底白字的风格不仅统一了视觉语言,更在复杂背景下提供了极佳的辨识度。

如果你觉得这个脚本有用,欢迎点赞、收藏,并关注我的 CSDN 账号,获取更多深度学习实用工具!

http://www.cnnetsun.cn/news/192285.html

相关文章:

  • 面向工业控制的模拟电子技术基础入门必看指南
  • 65、利用组策略管理硬件访问与打印机分配全攻略
  • f1系列替换下载失败
  • LangFlow内置模板库发布,涵盖常见AI应用场景
  • Centos7安装Maven环境
  • 【Arbess】1、安装Arbess
  • 实战案例:Arduino Uno R3开发板读取加速度传感器数据
  • 集体好奇心与团队学习能力的正相关
  • 树莓派5安装ROS2快速理解操作流程
  • LangFlow企业文化宣传文案生成工具
  • Java SpringBoot+Vue3+MyBatis 太原学院商铺管理系统系统源码|前后端分离+MySQL数据库
  • LangFlow员工满意度调查问卷生成器
  • SpringBoot+Vue 网上宠物店系统平台完整项目源码+SQL脚本+接口文档【Java Web毕设】
  • 公考日记9
  • Screen to Gif帧率调整的正确姿势
  • LangFlow走失老人定位协助流程设计
  • 操作系统移植视角下的x64和arm64差异:核心要点
  • vivo X300系列凭什么更受欢迎?旗舰体验这次更到位
  • 新手避坑指南:multisim14.3下载安装时防病毒误删技巧
  • LangFlow水族箱生态监控报警系统设想
  • elasticsearch可视化工具实现集群负载均衡监控教程
  • 自创的机械臂新算法,因为是AI写的,暂时,并不智能,但目前支持任何段数
  • OrCAD与Allegro协同工作:无缝对接设计流程
  • 从零实现无乱码开发环境:Keil + UTF-8-BOM配置教程
  • 调整IDE设置以避免代码自动换行
  • Java面试官怒怼水货程序员:Spring Cloud微服务+Kafka消息队列+Redis缓存,你到底会不会?
  • HBuilderX运行网页空白或报错?图解说明核心要点
  • Windows下Arduino安装教程:从下载到IDE配置手把手指导
  • 并网型直驱永磁同步风力发电系统simulink仿真
  • 如何为色盲人士创建可访问的图表