如何将Polygon格式转换为YOLO v5格式？代码转换坐标错误求助

问题描述 #

现有一个Polygon格式的边界框：

POLYGON ((799 1776, 799 2016, 490 2016, 490 1776, 799 1776))

需要转换为YOLO v5格式，但运行以下Python代码后得到错误坐标0 389.0 1552.0 160.0 165.0，不符合YOLO格式要求，需要修正代码。

原代码：

import logging
from pathlib import Path
import pandas as pd
from shapely.wkt import loads
yolo_output_dir = Path("my_yolo")
yolo_output_dir.mkdir(parents=True, exist_ok=True)
df = (pd
      .read_csv('images_bboxes.csv')
      .fillna(value={'geometry': '2'}))

worm_types = {
    y: x for (x, y) in enumerate(df['worm_type'].unique())
}
logging.critical(worm_types)

for (i, g) in df.groupby('image_id', sort=False):
    dst = yolo_output_dir.joinpath(i).with_suffix('.txt')
    logging.warning(dst)
    with dst.open('w') as fp:
        for i in g.itertuples(index=False):
            if i.geometry:
                worm = worm_types[i.worm_type]
                geometry = loads(i.geometry)
                (minx, miny, maxx, maxy) = geometry.bounds
                (w, h) = (maxx - minx, maxy - miny)
                print(worm, minx, miny, w, h, file=fp)

错误原因 #

YOLO v5的标注格式要求是归一化后的数值，格式为：

<类别ID> <归一化中心x> <归一化中心y> <归一化宽度> <归一化高度>

原代码存在两个核心错误：

• 直接输出了原始像素的左上角坐标（minx、miny），而非边界框的中心坐标
• 未将坐标和尺寸除以对应图像的宽高进行归一化，导致数值是原始像素值，不符合YOLO要求

修正方案 #

要完成正确转换，需要：

1. 获取每个图像的实际宽度和高度（可从CSV字段读取，或通过图像文件读取）
2. 计算边界框的中心坐标：center_x = (minx + maxx) / 2，center_y = (miny + maxy) / 2
3. 将中心坐标、宽、高分别除以图像宽、高，得到0-1之间的归一化值

修正后的代码 #

假设你的CSV文件中包含image_width和image_height字段，对应每个图像的尺寸；如果没有，可替换为通过PIL读取图像文件获取尺寸的逻辑：

import logging
from pathlib import Path
import pandas as pd
from shapely.wkt import loads
# 如果需要从图像文件读取尺寸，需导入PIL
# from PIL import Image

yolo_output_dir = Path("my_yolo")
yolo_output_dir.mkdir(parents=True, exist_ok=True)
df = (pd
      .read_csv('images_bboxes.csv')
      .fillna(value={'geometry': '2'}))

worm_types = {
    y: x for (x, y) in enumerate(df['worm_type'].unique())
}
logging.critical(worm_types)

for (image_id, g) in df.groupby('image_id', sort=False):
    dst = yolo_output_dir.joinpath(image_id).with_suffix('.txt')
    logging.warning(dst)
    # 方式1：从CSV读取图像尺寸（假设CSV有image_width和image_height列）
    img_width = g.iloc[0]['image_width']
    img_height = g.iloc[0]['image_height']
    
    # 方式2：如果CSV没有尺寸，从图像文件读取（需替换为你的图像路径）
    # img_path = Path("your_images_dir") / image_id
    # with Image.open(img_path) as img:
    #     img_width, img_height = img.size

    with dst.open('w') as fp:
        for row in g.itertuples(index=False):
            if row.geometry and row.geometry != '2':
                worm = worm_types[row.worm_type]
                geometry = loads(row.geometry)
                minx, miny, maxx, maxy = geometry.bounds
                # 计算中心坐标和宽高
                center_x = (minx + maxx) / 2
                center_y = (miny + maxy) / 2
                bbox_width = maxx - minx
                bbox_height = maxy - miny
                # 归一化处理
                norm_center_x = center_x / img_width
                norm_center_y = center_y / img_height
                norm_width = bbox_width / img_width
                norm_height = bbox_height / img_height
                # 写入文件，保留足够小数位
                print(f"{worm} {norm_center_x:.6f} {norm_center_y:.6f} {norm_width:.6f} {norm_height:.6f}", file=fp)

验证示例 #

针对给出的Polygon：

• 计算得minx=490, miny=1776, maxx=799, maxy=2016
• 中心x=(490+799)/2=644.5，中心y=(1776+2016)/2=1896
• 宽=799-490=309，高=2016-1776=240
  假设图像宽为1920，高为2560，归一化后的值为：
  0 644.5/1920≈0.3357 1896/2560≈0.7406 309/1920≈0.1609 240/2560≈0.09375
  输出格式为0 0.335729 0.740625 0.160938 0.093750，符合YOLO要求

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