🐛 fix: label studio to coco

Zerohertz · Nov 4, 2024 · 0a3a176 · 0a3a176
1 parent 40f45eb
commit 0a3a176
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 32 deletions.
diff --git a/test/test_vision.py b/test/test_vision.py
@@ -563,6 +563,10 @@ def test_LabelStudio_detection():
     assert "test.jpg" in os.listdir(os.path.join(yolo_path, "images"))
     assert "test.txt" in os.listdir(os.path.join(yolo_path, "labels"))
 
+    coco_path = "label-studio-detection-coco"
+    ls.coco(coco_path, {"Cat": 1})
+    assert f"{coco_path}.json" in os.listdir()
+
 
 def test_LabelStudio_segmentation():
     ls = zz.vision.LabelStudio(
@@ -588,6 +592,10 @@ def test_LabelStudio_segmentation():
     assert "test.jpg" in os.listdir(os.path.join(yolo_path, "images"))
     assert "test.txt" in os.listdir(os.path.join(yolo_path, "labels"))
 
+    coco_path = "label-studio-segmentation-coco"
+    ls.coco(coco_path, {"Cat": 1})
+    assert f"{coco_path}.json" in os.listdir()
+
 
 def _test_YoloLoader_detection(path=None):
     vis_path = "yololoader-detection"
@@ -673,32 +681,6 @@ def test_CocoLoader_segmentation():
 
 
 """
-TODO: LabelStudio.coco()
->>> lss.coco({"Cat": 1})
-100%|███████████████████████████████████████████████████| 1/1 [00:00<00:00, 5769.33it/s]
-Traceback (most recent call last):
-  File "<stdin>", line 1, in <module>
-  File "/home/zerohertz/Zerohertz/zerohertzLib/zerohertzLib/vision/data.py", line 541, i
-n coco
-    write_json(converted_gt, self.data_path)
-  File "/home/zerohertz/Zerohertz/zerohertzLib/zerohertzLib/util/json.py", line 362, in 
-write_json
-    file.write(orjson.dumps(data, option=orjson.OPT_INDENT_2))
-               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-TypeError: Type is not JSON serializable: numpy.float64
->>> lsd.coco({"Cat": 1})
-100%|███████████████████████████████████████████████████| 1/1 [00:00<00:00, 8112.77it/s]
-Traceback (most recent call last):
-  File "<stdin>", line 1, in <module>
-  File "/home/zerohertz/Zerohertz/zerohertzLib/zerohertzLib/vision/data.py", line 541, i
-n coco
-    write_json(converted_gt, self.data_path)
-  File "/home/zerohertz/Zerohertz/zerohertzLib/zerohertzLib/util/json.py", line 362, in 
-write_json
-    file.write(orjson.dumps(data, option=orjson.OPT_INDENT_2))
-               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-TypeError: Type is not JSON serializable: numpy.float64
-
 TODO: YoloLoader.labelstudio()
 >>> yolo.labelstudio()
 Traceback (most recent call last):

diff --git a/zerohertzLib/vision/data.py b/zerohertzLib/vision/data.py
@@ -478,19 +478,20 @@ def classification(
                             f"Impossible crop ('x_0': {x_0}, 'y_0': {y_0}, 'x_1': {x_1}, 'y_1': {y_1})"
                         )
 
-    def coco(self, label: Dict[str, int]) -> None:
+    def coco(self, target_path: str, label: Dict[str, int]) -> None:
         """Label Studio로 annotation한 JSON data를 COCO format으로 변환
 
         Args:
+            target_path (``str``): COCO format data가 저장될 경로
             label (``Optional[Dict[str, int]]``): Label Studio에서 사용한 label을 변경하는 dictionary
 
         Returns:
-            ``None``: ``{data_path}.json`` 에 JSON file 저장
+            ``None``: ``{target_path}.json`` 에 JSON file 저장
 
         Examples:
             >>> ls = zz.vision.LabelStudio(data_path, json_path)
             >>> label = {"label1": 1, "label2": 2}
-            >>> ls.coco(label)
+            >>> ls.coco(target_path, label)
             100%|█████████████| 476/476 [00:00<00:00, 78794.25it/s]
         """
         converted_gt = {
@@ -524,18 +525,19 @@ def coco(self, label: Dict[str, int]) -> None:
                     box_cwh[:2] -= box_cwh[2:] / 2
                 else:
                     raise ValueError(f"Unknown annotation type: {self.type}")
+                box_cwh = box_cwh.tolist()
                 _annotations.append(
                     {
                         "segmentation": [poly.reshape(-1).tolist()],
                         "area": box_cwh[2] * box_cwh[3],
                         "iscrowd": 0,
-                        "image_id": id,
-                        "bbox": box_cwh.tolist(),
+                        "image_id": id_,
+                        "bbox": box_cwh,
                         "category_id": label[lab],
                         "id": ant_id + ant_id_,
                     }
                 )
             converted_gt["images"].append(_images)
             converted_gt["annotations"] += _annotations
             ant_id += len(result["labels"]) + 1
-        write_json(converted_gt, self.data_path)
+        write_json(converted_gt, target_path)