Merge pull request #9 from hannarud/patch-1

Support for any category declared in CATEGORIES
waspinator · Jul 24, 2018 · d29534e · d29534e
2 parents db5efbf + 8d30e18
commit d29534e
Showing 1 changed file with 1 addition and 6 deletions.
diff --git 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)