Generate yaml including class id

Author: darouwan

general_json2yolo.py

@@ -1,6 +1,6 @@
 import contextlib
 import json
-
+import yaml
 import cv2
 import pandas as pd
 from PIL import Image
@@ -254,13 +254,18 @@ def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91
     save_dir = make_dirs()  # output directory
     coco80 = coco91_to_coco80_class()
 
+    categories_dict = {}
+
     # Import json
     for json_file in sorted(Path(json_dir).resolve().glob('*.json')):
         fn = Path(save_dir) / 'labels' / json_file.stem.replace('instances_', '')  # folder name
         fn.mkdir()
         with open(json_file) as f:
             data = json.load(f)
 
+        # Create categories_dict
+        categories_dict = {item['id']: item['name'] for item in data['categories']}
+
         # Create image dict
         images = {'%g' % x['id']: x for x in data['images']}
         # Create image-annotations dict
@@ -308,6 +313,16 @@ def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91
                     line = *(segments[i] if use_segments else bboxes[i]),  # cls, box or segments
                     file.write(('%g ' * len(line)).rstrip() % line + '\n')
 
+        yaml_dict = {
+            "names": {v: k for k, v in categories_dict.items()},
+            # "path": "yolo_datasets",
+            "train": "images/train",
+            "val": "images/val"
+        }
+
+        with open((Path(save_dir) / json_file.stem).with_suffix('.yaml'), "w") as f:
+            yaml.dump(yaml_dict, f)
+
 
 def min_index(arr1, arr2):
     """Find a pair of indexes with the shortest distance.