bdd100k_to_YOLO/bdd_to_yolo_converter.py at master · RocketFlash/bdd100k_to_YOLO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import json
import os
from tqdm import tqdm
from config import PATHS, categories, img_size


def generate_yolo_labels(json_path, save_path, fname_prefix=None, fname_postfix=None):
    img_w, img_h = img_size
    ignore_categories = ["drivable area", "lane"]
    with open(json_path) as json_file:
        data = json.load(json_file)
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        for img in tqdm(data):
            img_name = str(img['name'][:-4])
            img_label_txt = (
                fname_prefix if fname_prefix is not None else '') + img_name
            img_label_txt += (
                fname_postfix if fname_postfix is not None else '') + ".txt"
            img_labels = [l for l in img['labels']
                          if l['category'] not in ignore_categories]
            file_path = save_path + img_label_txt
            with open(file_path, 'w+') as f_label:
                for label in img_labels:
                    y1 = label['box2d']['y1']
                    x2 = label['box2d']['x2']
                    x1 = label['box2d']['x1']
                    y2 = label['box2d']['y2']
                    class_name = label['category']
                    class_id = categories[class_name]

                    bbox_x = (x1 + x2)/2
                    bbox_y = (y1 + y2)/2

                    bbox_width = x2-x1
                    bbox_height = y2-y1

                    bbox_x_norm = bbox_x / img_w
                    bbox_y_norm = bbox_y / img_h

                    bbox_width_norm = bbox_width / img_w
                    bbox_height_norm = bbox_height / img_h

                    line_to_write = '{} {} {} {} {}'.format(
                        class_id, bbox_x_norm, bbox_y_norm, bbox_width_norm, bbox_height_norm)
                    f_label.write(line_to_write + "\n")


def generate_yolo_filenames(imgs_path, labels_path, output_file_path):
    os.makedirs(PATHS['save_path'] , exist_ok=True)
    count = 0
    with open(output_file_path, 'w+') as f:
        for dirpath, dirs, files in os.walk(imgs_path):
            for filename in tqdm(files):
                is_exist = os.path.isfile(labels_path+filename[:-4] + '.txt')
                if is_exist:
                    if filename.endswith(".jpg"):
                        file_path = os.path.join(
                            PATHS['owner_prefix'], dirpath, filename)
                        f.write(file_path + "\n")
                else:
                    count += 1
    print('number of skipped files: {}'.format(count))


def generate_names_file(filename='bdd100k.names'):
    os.makedirs(PATHS['save_path'] , exist_ok=True)
    output_file_path = PATHS['save_path'] + filename
    with open(output_file_path, 'w+') as f:
        for key, _ in categories.items():
            f.write(key + "\n")
    print('{} created'.format(filename))


def generate_data_file(filename='bdd100k.data'):
    os.makedirs(PATHS['save_path'] , exist_ok=True)
    output_file_path = PATHS['save_path'] + filename
    with open(output_file_path, 'w+') as f:
        f.write('classes = 10' + '\n')
        f.write('train = {}train.txt'.format(PATHS['save_path']) + '\n')
        f.write('valid = {}val.txt'.format(PATHS['save_path']) + '\n')
        f.write('test = {}test.txt'.format(PATHS['save_path']) + '\n')
        f.write('names = {}bdd100k.names'.format(PATHS['save_path']) + '\n')
        f.write('backup = {}backup'.format(PATHS['save_path']) + '\n')
    print('{} created'.format(filename))


if __name__ == '__main__':

    print('start YOLO labels creation')
    generate_yolo_labels(PATHS['labels_path_json_train'], PATHS['labels_save_path_train'],
                         fname_prefix='train_', fname_postfix=None)
    generate_yolo_labels(PATHS['labels_path_json_val'], PATHS['labels_save_path_val'],
                         fname_prefix='val_', fname_postfix=None)

    # # generate labels for augmented images
    # generate_yolo_labels(PATHS['labels_path_json_train'], PATHS['labels_save_path_train'],
    #                      fname_prefix='train_', fname_postfix='_fake_B')
    # generate_yolo_labels(PATHS['labels_path_json_val'], PATHS['labels_save_path_val'],
    #                      fname_prefix='val_', fname_postfix='_fake_B')

    # # GENERATE TEST FILES
    generate_yolo_labels(PATHS['labels_path_json_train'], PATHS['labels_save_path_test'],
                         fname_prefix='train_', fname_postfix=None)
    generate_yolo_labels(PATHS['labels_path_json_val'], PATHS['labels_save_path_test'],
                         fname_prefix='val_', fname_postfix=None)

    print('start YOLO filenames file creation')
    generate_yolo_filenames(
        PATHS['images_path_train'], PATHS['labels_save_path_train'], PATHS['file_path_train'])
    generate_yolo_filenames(
        PATHS['images_path_test'], PATHS['labels_save_path_test'], PATHS['file_path_test'])


    generate_names_file(filename='bdd100k.names')
    generate_data_file(filename='bdd100k.data')