YOLOv5实战:如何一键导出检测框的坐标、类别和置信度到TXT文件(附完整代码)
YOLOv5实战结构化导出检测结果的工程化解决方案在计算机视觉项目的实际落地过程中我们常常需要将模型检测结果以结构化形式保存用于后续的数据分析、系统集成或模型评估。本文将深入探讨如何通过YOLOv5高效导出检测框的坐标、类别和置信度信息到TXT文件并提供可直接复用的工程化代码方案。1. 理解YOLOv5输出数据结构YOLOv5的检测结果本质上包含三个核心维度信息边界框坐标通常以(x_min, y_min, x_max, y_max)格式表示类别索引对应data/yaml文件中定义的类别顺序置信度分数模型对当前预测结果的把握程度0-1之间在原始实现中这些信息被封装在det变量中每个检测结果以[x1, y1, x2, y2, conf, cls]的形式存在。理解这一数据结构是进行定制化导出的基础。注意YOLOv5默认使用归一化坐标0-1之间但在实际应用中我们可能需要转换为绝对像素坐标。2. 核心导出方案设计与实现2.1 基础导出功能实现以下代码展示了如何修改YOLOv5的detect.py以实现结构化导出def write_detection_results(det, img_shape, txt_path): 将检测结果写入TXT文件 :param det: 检测结果张量 [n,6] (xyxy, conf, cls) :param img_shape: 原始图像尺寸 (height, width) :param txt_path: 输出文件路径 with open(txt_path, a) as f: for *xyxy, conf, cls in reversed(det): # 转换为绝对坐标 x1, y1, x2, y2 [int(coord) for coord in xyxy] # 格式化写入类别 置信度 x1 y1 x2 y2 line f{int(cls)} {conf:.4f} {x1} {y1} {x2} {y2}\n f.write(line)将此函数集成到detect.py的适当位置通常在绘制检测框之前并通过以下方式调用if save_txt: write_detection_results(det, im0.shape[:2], txt_path)2.2 高级导出格式定制根据不同的下游应用需求我们可以灵活调整输出格式。以下是几种常见变体格式类型数据结构适用场景基础格式cls conf x1 y1 x2 y2简单数据记录COCO风格cls conf x1 y1 width height兼容COCO评估工具归一化格式cls conf x_center y_center w hYOLO训练格式实现COCO风格导出的代码示例def write_coco_style(det, txt_path): with open(txt_path, a) as f: for *xyxy, conf, cls in reversed(det): x1, y1, x2, y2 xyxy width x2 - x1 height y2 - y1 line f{int(cls)} {conf:.4f} {x1} {y1} {width} {height}\n f.write(line)3. 工程化增强方案3.1 多文件批量处理在实际项目中我们常需要处理大量图像。以下脚本展示了如何批量处理并保持文件名对应import os from pathlib import Path def process_directory(source_dir, output_dir): source_dir Path(source_dir) output_dir Path(output_dir) output_dir.mkdir(exist_okTrue) # 获取所有图像文件 image_files list(source_dir.glob(*.jpg)) list(source_dir.glob(*.png)) for img_path in image_files: # 运行检测并导出结果 run_detection( weightsyolov5s.pt, sourcestr(img_path), projectstr(output_dir), nameresults, save_txtTrue ) # 重命名结果文件以匹配原图名 result_file output_dir / results / labels / f{img_path.stem}.txt if result_file.exists(): new_path output_dir / f{img_path.stem}_det.txt result_file.rename(new_path)3.2 性能优化技巧处理大规模数据时IO操作可能成为瓶颈。以下方法可显著提升导出效率缓冲写入累积一定数量的检测结果后批量写入并行处理使用多进程处理不同图像二进制格式对于极大数据集考虑使用二进制格式存储优化后的写入函数示例def buffered_write(detections, file_path, buffer_size100): buffer [] for det in detections: buffer.append(f{det[cls]} {det[conf]} {det[x1]} {det[y1]} {det[x2]} {det[y2]}) if len(buffer) buffer_size: with open(file_path, a) as f: f.write(\n.join(buffer) \n) buffer.clear() if buffer: # 写入剩余数据 with open(file_path, a) as f: f.write(\n.join(buffer) \n)4. 结果分析与可视化4.1 数据统计与质量检查导出数据后我们通常需要进行基本统计分析import pandas as pd import matplotlib.pyplot as plt def analyze_results(result_file): # 读取导出的TXT文件 df pd.read_csv(result_file, sep , names[class, confidence, x1, y1, x2, y2]) # 基本统计 print(f总检测数: {len(df)}) print(f平均置信度: {df[confidence].mean():.3f}) # 类别分布可视化 class_counts df[class].value_counts() class_counts.plot(kindbar) plt.title(Class Distribution) plt.xlabel(Class ID) plt.ylabel(Count) plt.show()4.2 结果可视化验证为确保导出数据的准确性建议实现简单的可视化验证import cv2 def visualize_detections(image_path, result_path): image cv2.imread(image_path) with open(result_path) as f: for line in f: cls, conf, x1, y1, x2, y2 map(float, line.strip().split()) # 绘制边界框 cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0,255,0), 2) # 添加标签 label f{int(cls)} {conf:.2f} cv2.putText(image, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2) cv2.imshow(Verification, image) cv2.waitKey(0)5. 高级应用场景5.1 与数据库集成对于需要长期存储和分析的场景可将结果直接导入数据库import sqlite3 def save_to_database(result_file, db_path): conn sqlite3.connect(db_path) cursor conn.cursor() # 创建表 cursor.execute( CREATE TABLE IF NOT EXISTS detections ( id INTEGER PRIMARY KEY AUTOINCREMENT, image_name TEXT, class INTEGER, confidence REAL, x1 INTEGER, y1 INTEGER, x2 INTEGER, y2 INTEGER, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP ) ) # 导入数据 with open(result_file) as f: for line in f: cls, conf, x1, y1, x2, y2 line.strip().split() cursor.execute( INSERT INTO detections (image_name, class, confidence, x1, y1, x2, y2) VALUES (?, ?, ?, ?, ?, ?, ?), (Path(result_file).stem, cls, conf, x1, y1, x2, y2) ) conn.commit() conn.close()5.2 自动化报告生成结合检测结果生成PDF报告from fpdf import FPDF def generate_report(result_file, output_pdf): pdf FPDF() pdf.add_page() pdf.set_font(Arial, size12) # 添加标题 pdf.cell(200, 10, txtDetection Results Report, ln1, alignC) # 添加统计信息 with open(result_file) as f: lines f.readlines() pdf.cell(200, 10, txtfTotal Detections: {len(lines)}, ln1) # 添加样本数据 pdf.cell(200, 10, txtSample Detections:, ln1) for line in lines[:10]: # 显示前10个检测结果 pdf.cell(200, 10, txtline.strip(), ln1) pdf.output(output_pdf)在实际项目中这种结构化导出方案显著提升了后续处理效率。一个典型的应用场景是将检测结果导入到GIS系统进行空间分析此时精确的坐标信息和规范的格式至关重要。

相关新闻

最新新闻

日新闻

周新闻

月新闻