dataset_scripts/crop_boxes.py at master · andrey1908/dataset_scripts · 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
import argparse
import json
import os
from PIL import Image
from tqdm import tqdm


def build_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('-json', '--json-file', required=True, type=str)
    parser.add_argument('-img-fld', '--images-folder', required=True, type=str)
    parser.add_argument('-cls', '--classes', nargs='+', type=str)
    parser.add_argument('-out-fld', '--out-folder', required=True, type=str)
    parser.add_argument('-log', '--log-file', type=str)
    return parser


def get_image_id_to_annotations_idxs(images, annotations):
    image_id_to_annotations_idxs = dict()
    for image in images:
        image_id_to_annotations_idxs[image['id']] = list()
    for i, ann in enumerate(annotations):
        image_id_to_annotations_idxs[ann['image_id']].append(i)
    return image_id_to_annotations_idxs


def get_category_id_to_name(categories):
    category_id_to_name = dict()
    for category in categories:
        category_id_to_name[category['id']] = category['name']
    return category_id_to_name


def set_between(x, a, b):
    da = x - a
    db = x - b
    if da * db <= 0:
        return x
    if abs(da) < abs(db):
        return a
    else:
        return b


def preprocess_box(box, im_w, im_h):
    box[2] += box[0]
    box[3] += box[1]
    box[0] = set_between(box[0], 0, im_w)
    box[1] = set_between(box[1], 0, im_h)
    box[2] = set_between(box[2], 0, im_w)
    box[3] = set_between(box[3], 0, im_h)
    box[:] = [round(b) for b in box]


def crop_boxes(json_dict, images_folder, classes, out_folder):
    log = list()
    image_id_to_annotations_idxs = get_image_id_to_annotations_idxs(json_dict['images'], json_dict['annotations'])
    category_id_to_name = get_category_id_to_name(json_dict['categories'])
    boxes_counters = dict()
    fill_classes = False
    if classes == None:
        classes = list()
        fill_classes = True
    for category in json_dict['categories']:
        if fill_classes:
            boxes_counters[category['id']] = 1
            os.makedirs(os.path.join(out_folder, category['name']), exist_ok=True)
            classes.append(category['name'])
        elif category['name'] in classes:
            boxes_counters[category['id']] = 1
            os.makedirs(os.path.join(out_folder, category['name']), exist_ok=True)

    for image in tqdm(json_dict['images']):
        log_image = {'original_file_name': image['file_name'], 'crops': list()}
        im = Image.open(os.path.join(images_folder, image['file_name']))
        annotations_idxs = image_id_to_annotations_idxs[image['id']]
        for annotation_idx in annotations_idxs:
            ann = json_dict['annotations'][annotation_idx]
            cat_id = ann['category_id']
            cat_name = category_id_to_name[cat_id]
            if cat_name not in classes:
                continue
            box = list(ann['bbox'])
            preprocess_box(box, im.size[0], im.size[1])
            if (box[0] >= box[2]) or (box[1] >= box[3]):
                continue
            croped = im.crop(tuple(box))
            croped.save(os.path.join(out_folder, cat_name, '{}.png'.format(boxes_counters[cat_id])))
            log_crop = {'cropped_file_name': os.path.join(cat_name, '{}.png'.format(boxes_counters[cat_id])),
                        'box': box}  # xtl, ytl, xbr, ybr
            log_image['crops'].append(log_crop)
            boxes_counters[cat_id] += 1
        log.append(log_image)
    return log


if __name__ == '__main__':
    parser = build_parser()
    args = parser.parse_args()
    with open(args.json_file, 'r') as f:
        json_dict = json.load(f)
    log = crop_boxes(json_dict, args.images_folder, args.classes, args.out_folder)
    if args.log_file is not None:
        with open(args.log_file, 'w') as f:
            json.dump(log, f, indent = 2)