46 lines
1.3 KiB
Python
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)
|
|
|