如何在YOLO-V8实例分割模型预测时将边界框值加入标签文本文件？

方案1：重新预测时生成带边界框的标签 #

若尚未完成预测，可通过Python脚本调用YOLOv8，自定义输出格式以同时包含边界框和分割多边形：

from ultralytics import YOLO
import os

# 加载训练好的模型
model = YOLO('/weights/best.pt')
# 执行预测
results = model.predict(source='/test/images', conf=0.25, imgsz=1024, save=True)

# 遍历结果，生成自定义标签
save_dir = results[0].save_dir
for result in results:
    img_base = os.path.splitext(os.path.basename(result.path))[0]
    label_path = os.path.join(save_dir, 'labels', f'{img_base}.txt')
    with open(label_path, 'w') as f:
        for cls, bbox, seg, conf in zip(result.boxes.cls, result.boxes.xywhn, result.masks.xy, result.boxes.conf):
            # 转换为字符串格式
            bbox_str = ' '.join(map(str, bbox.tolist()))
            seg_str = ' '.join(map(str, seg.flatten().tolist()))
            # 写入行：类别 + 归一化边界框 + 多边形坐标 + 置信度
            f.write(f'{int(cls)} {bbox_str} {seg_str} {conf.item()}\n')

方案2：批量处理已生成的标签文件 #

若已通过CLI完成预测，可运行以下脚本从多边形坐标计算边界框，并添加到现有标签中：

import os

def get_normalized_bbox(seg_coords):
    # 从多边形坐标中提取所有x、y值
    xs = seg_coords[::2]
    ys = seg_coords[1::2]
    x_min, x_max = min(xs), max(xs)
    y_min, y_max = min(ys), max(ys)
    # 计算归一化边界框参数
    x_center = (x_min + x_max) / 2
    y_center = (y_min + y_max) / 2
    width = x_max - x_min
    height = y_max - y_min
    return [x_center, y_center, width, height]

# 替换为你的标签目录路径（默认是runs/segment/predict/labels）
LABELS_DIR = '/your/labels/directory/path'

for filename in os.listdir(LABELS_DIR):
    if not filename.endswith('.txt'):
        continue
    file_path = os.path.join(LABELS_DIR, filename)
    
    with open(file_path, 'r') as f:
        lines = f.read().splitlines()
    
    updated_lines = []
    for line in lines:
        if not line:
            continue
        parts = list(map(float, line.split()))
        cls = int(parts[0])
        seg_coords = parts[1:-1]  # 类别后、置信度前的部分是多边形坐标
        conf = parts[-1]
        
        # 计算边界框
        bbox = get_normalized_bbox(seg_coords)
        # 拼接新行
        updated_line = f'{cls} {" ".join(map(str, bbox))} {" ".join(map(str, seg_coords))} {conf}'
        updated_lines.append(updated_line)
    
    # 覆盖原文件（如需保留原文件，可修改为写入新目录）
    with open(file_path, 'w') as f:
        f.write('\n'.join(updated_lines))

说明 #

• 边界框计算逻辑：从分割多边形的所有坐标中找到x/y的极值，推导归一化的中心坐标、宽度和高度
• 处理后标签格式：<类别索引> <x-center> <y-center> <width> <height> <多边形x1> <多边形y1> ... <多边形xn> <多边形yn> <置信度>

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