CODEWITHSUNDEEP
pythonmachine-learningneural-networkyolov5yolov4
akshath mugad
akshath mugad

Reputation: 21

How to convert Polygon format to YOLO forma

POLYGON ((799 1776, 799 2016, 490 2016, 490 1776, 799 1776)) This is the bounding box in POLYGON i want this in YOLO v5 format

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)

This is the code i tried,, but it is giving wrong coordinates..

POLYGON ((799 1776, 799 2016, 490 2016, 490 1776, 799 1776)) is converted to 0 389.0 1552.0 160.0 165.0 which is wrong

Upvotes: 2

Views: 429

Answers (1)

JNM
JNM

Reputation: 118

I used this function to convert a WKT polygon into x, y, w, h coordinates that can be used to train a YOLO model.

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

Upvotes: 0

Related Questions