news-navigator/change-to-yolo-format.py
2022-08-02 10:09:06 +00:00

46 lines
1.3 KiB
Python

# start od 0
import cv2
from labels_list import labels_list
from tqdm import tqdm
import sys
dataset = sys.argv[1]
def to_yolo_format(entity, y_total, x_total):
label_str, ref = entity.split(':')
label_id = labels_list.index(label_str)
x0,y0,x1,y1 = [float(i) for i in ref.split(',')]
x_middle = (x0 + x1) / 2
y_middle = (y0 + y1) / 2
x_w = x1 - x0
y_w = y1 - y0
x_middle /= x_total
y_middle /= y_total
x_w /= x_total
y_w /= y_total
items_to_write = (label_id, x_middle, y_middle, x_w, y_w)
items_to_write = [str(i) for i in items_to_write]
str_to_write = ' '.join(items_to_write) + '\n'
return str_to_write
with open(f'{dataset}/expected.tsv','r') as f_exp, open(f'{dataset}/in.tsv') as f_in:
for line_in, line_exp in zip(f_in, f_exp):
img_name = line_in.rstrip()
img_references = line_exp.rstrip().split(' ')
img = cv2.imread('images/' + img_name)
img_dimensions = img.shape
y_total, x_total, _ = img_dimensions
with open(f'{dataset}-dataset/labels/' + img_name.replace('.jpg','.txt'), 'w') as f_out:
for entity in img_references:
str_to_write = to_yolo_format(entity, y_total, x_total)
f_out.write(str_to_write)