From 21443f11d91b0108ca4e446314b444face98fa41 Mon Sep 17 00:00:00 2001
From: JF Chen <k-2feng@hotmail.com>
Date: Tue, 26 Mar 2024 15:35:30 +0800
Subject: [PATCH] Merge modification

---
 general_json2yolo.py | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/general_json2yolo.py b/general_json2yolo.py
index f1ada2c..e01e151 100644
--- a/general_json2yolo.py
+++ b/general_json2yolo.py
@@ -1,6 +1,8 @@
 import contextlib
 import json
+import shutil
 
+import yaml
 import cv2
 import pandas as pd
 from PIL import Image
@@ -254,13 +256,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
+    for json_file in sorted(Path(json_dir).resolve().glob('*coco*.json')):
+        fn = Path(save_dir) / 'labels' / json_file.stem.replace('instances_', '').replace('_coco','')  # folder name
         fn.mkdir()
         with open(json_file) as f:
             data = json.load(f)
 
+        # Create categories_dict
+        categories_dict = {item['id'] - 1: 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,9 +315,20 @@ 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": {k: v 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)
+        shutil.copy((Path(save_dir) / json_file.stem).with_suffix('.yaml'), (Path(save_dir) / 'classes').with_suffix('.yaml'))
+
 
 def min_index(arr1, arr2):
-    """Find a pair of indexes with the shortest distance. 
+    """Find a pair of indexes with the shortest distance.
     Args:
         arr1: (N, 2).
         arr2: (M, 2).
@@ -324,12 +342,12 @@ def min_index(arr1, arr2):
 def merge_multi_segment(segments):
     """Merge multi segments to one list.
     Find the coordinates with min distance between each segment,
-    then connect these coordinates with one thin line to merge all 
+    then connect these coordinates with one thin line to merge all
     segments into one.
 
     Args:
         segments(List(List)): original segmentations in coco's json file.
-            like [segmentation1, segmentation2,...], 
+            like [segmentation1, segmentation2,...],
             each segmentation is a list of coordinates.
     """
     s = []