🔀 [Merge] branch 'DATASET' into TEST

WongKinYiu · Jul 24, 2024 · 868c821 · 868c821
2 parents 9c96910 + d2d621f
commit 868c821
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/yolo/utils/dataset_utils.py b/yolo/utils/dataset_utils.py
@@ -100,7 +100,10 @@ def scale_segmentation(
     h, w = image_dimensions["height"], image_dimensions["width"]
     for anno in annotations:
         category_id = anno["category_id"]
-        seg_list = [item for sublist in anno["segmentation"] for item in sublist]
+        if "segmentation" in anno:
+            seg_list = [item for sublist in anno["segmentation"] for item in sublist]
+        elif "bbox" in anno:
+            seg_list = anno["bbox"]
         scaled_seg_data = (
             np.array(seg_list).reshape(-1, 2) / [w, h]
         ).tolist()  # make the list group in x, y pairs and scaled with image width, height