From 8d30e1845c4f136b1830edac11413a234a137bef Mon Sep 17 00:00:00 2001 From: Hanna Rudakouskaya Date: Tue, 24 Jul 2018 18:11:22 +0300 Subject: [PATCH] Support for any category declared in CATEGORIES --- examples/shapes/shapes_to_coco.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/examples/shapes/shapes_to_coco.py b/examples/shapes/shapes_to_coco.py index 75434a2..2942ceb 100755 --- a/examples/shapes/shapes_to_coco.py +++ b/examples/shapes/shapes_to_coco.py @@ -99,12 +99,7 @@ def main(): for annotation_filename in annotation_files: print(annotation_filename) - if 'square' in annotation_filename: - class_id = 1 - elif 'circle' in annotation_filename: - class_id = 2 - else: - class_id = 3 + class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0] category_info = {'id': class_id, 'is_crowd': 'crowd' in image_filename} binary_mask = np.asarray(Image.open(annotation_filename)