yolo_datagen.py

import os
from pinpoint.jdf import JDFDataset, vis

dataset_path = '.../data'
num_epochs_gen = 3
size = (224, 224)
box_size = (10, 10)

output_yolo_dataset_path = os.path.join(os.path.dirname(__file__), 'yolo_dataset')
if not os.path.exists(output_yolo_dataset_path):
    os.mkdir(output_yolo_dataset_path)

dataset = JDFDataset(dataset_path, img_reshape=size, coordconv=False, normalize=False)

idx = 0
for epoch in range(num_epochs_gen):
    for img, label in dataset:
        imgstr = ''
        vis(img, duration=0, denormalize=False, savePath=output_yolo_dataset_path)
        for c, x, y in label:
            if c == -1:
                break
            imgstr += "{} {} {} {} {}\n".format(c, x, y, box_size[0]/size[0], box_size[1]/size[1])
        # print(imgstr)

        n = sum(1 for f in os.listdir(output_yolo_dataset_path) if os.path.isfile(os.path.join(output_yolo_dataset_path, f)))
        with open(os.path.join(output_yolo_dataset_path, "Figure_{}.txt".format(n)), 'w') as f:
            f.write(imgstr)