[Fix] Fix bugs in the output annotation file when --class-id-txt is…

… set. (#430) * [Fix] Fix bugs in the output annotation file when is set. * Update labelme2coco.py * Update labelme2coco.py * Improve code Co-authored-by: HinGwenWoong <[email protected]>
open-mmlab · Jan 6, 2023 · 2290713 · 2290713
1 parent d8cbf21
commit 2290713
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/tools/dataset_converters/labelme2coco.py b/tools/dataset_converters/labelme2coco.py
@@ -171,6 +171,13 @@ def parse_labelme_to_coco(
         category_to_id = all_classes_id
         categories_labels = list(all_classes_id.keys())
 
+        # add class_ids and class_names to the categories list in coco_json
+        for class_name, class_id in category_to_id.items():
+            coco_json['categories'].append({
+                'id': class_id,
+                'name': class_name
+            })
+
     # filter incorrect image file
     img_file_list = [
         img_file for img_file in Path(image_dir).iterdir()
@@ -283,7 +290,7 @@ def convert_labelme_to_coco(image_dir: str,
                                  ' words, like "1 Big house" -> "1 '
                                  'Big-house".')
             v, k = class_info
-            all_classes_id.update({k: v})
+            all_classes_id.update({k: int(v)})
     else:
         all_classes_id = None