如何将Polygon格式转换为YOLO格式

Question

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

这是 POLYGON 中的边界框 我想要 YOLO v5 格式的

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)

这是我尝试过的代码，但它给出了错误的坐标..

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

被转换为

0 389.0 1552.0 160.0 165.0

这是错误的

Answer 1

我使用此函数将 WKT 多边形转换为可用于训练 YOLO 模型的 x、y、w、h 坐标。

import shapely.wkt

def polygon_to_yolo(polygon, img_width=100, img_height=100):
    """Convert string defining WKT polygon into YOLO format"""
    p = shapely.wkt.loads(polygon)
    x, y = p.centroid.x, p.centroid.y
    w, h = p.bounds[2] - p.bounds[0], p.bounds[3] - p.bounds[1]
    return x / img_width, y / img_height, w / img_width, h / img_height