fix: FIT-1653: COCO Export Category IDs mismatch

GitOrigin-RevId: a0484ba33aeaf90344482fbbe2317eb7ee7f3294

Author: matt-bernstein

src/label_studio_sdk/converter/converter.py

@@ -1,4 +1,5 @@
 import io
+import itertools
 import logging
 import math
 import os
@@ -1250,20 +1251,25 @@ def _get_labels(self):
             labels |= set(info["labels"])
             attrs = info["labels_attrs"]
             for label in attrs:
-                if attrs[label].get("category"):
-                    categories.append(
-                        {"id": attrs[label].get("category"), "name": label}
+                raw_category_id = attrs[label].get("category")
+                if raw_category_id is None:
+                    continue
+                try:
+                    category_id = int(raw_category_id)
+                except (TypeError, ValueError):
+                    logger.warning(
+                        "Invalid category id %r for label %r in tag %r, skipping explicit category mapping",
+                        raw_category_id,
+                        label,
+                        name,
                     )
-                    category_name_to_id[label] = attrs[label].get("category")
+                    continue
+                categories.append({"id": category_id, "name": label})
+                category_name_to_id[label] = category_id
         labels_to_add = set(labels) - set(list(category_name_to_id.keys()))
-        labels_to_add = sorted(list(labels_to_add))
-        idx = 0
-        while idx in list(category_name_to_id.values()):
-            idx += 1
-        for label in labels_to_add:
+        used_category_ids = set(category_name_to_id.values())
+        idxs = (c for c in itertools.count() if c not in used_category_ids)
+        for label, idx in zip(sorted(labels_to_add), idxs):
             categories.append({"id": idx, "name": label})
             category_name_to_id[label] = idx
-            idx += 1
-            while idx in list(category_name_to_id.values()):
-                idx += 1
         return categories, category_name_to_id

src/label_studio_sdk/converter/keypoints.py

@@ -21,29 +21,30 @@ def update_categories_for_keypoints(categories, category_name_to_id, label_confi
             keypoint_labels.extend(cfg.get("labels", []))
     keypoint_labels = list(dict.fromkeys(keypoint_labels))
 
+    if not keypoint_labels:
+        return categories, category_name_to_id
+
     non_kp = [cat.copy() for cat in categories if cat["name"] not in keypoint_labels]
 
     new_categories = []
     new_mapping = {}
-    next_id = 0
     for cat in non_kp:
-        cat["id"] = next_id
+        category_id = cat["id"]
         new_categories.append(cat)
-        new_mapping[cat["name"]] = next_id
-        next_id += 1
-
-    if keypoint_labels:
-        merged_id = next_id
-        merged_category = {
-            "id": merged_id,
-            "name": "default",
-            "supercategory": "default",
-            "keypoints": keypoint_labels,
-            "skeleton": []
-        }
-        new_categories.append(merged_category)
-        for kp_name in keypoint_labels:
-            new_mapping[kp_name] = merged_id
+        new_mapping[cat["name"]] = category_id
+
+    numeric_ids = [cat["id"] for cat in new_categories if isinstance(cat["id"], int)]
+    merged_id = (max(numeric_ids) + 1) if numeric_ids else 0
+    merged_category = {
+        "id": merged_id,
+        "name": "default",
+        "supercategory": "default",
+        "keypoints": keypoint_labels,
+        "skeleton": []
+    }
+    new_categories.append(merged_category)
+    for kp_name in keypoint_labels:
+        new_mapping[kp_name] = merged_id
 
     return new_categories, new_mapping
 

tests/custom/converter/test_export_coco.py

@@ -159,3 +159,79 @@ def test_convert_to_coco_keypoints(temp_out_dir: Path):
     cat = kp_cats[0]
 
     assert cat["name"] == "default"
+
+
+def test_convert_to_coco_preserves_explicit_category_ids(temp_out_dir: Path):
+    """COCO export should keep explicit `category` ids from the labeling config."""
+    config = """
+    <View>
+      <Image name="image" value="$image"/>
+      <RectangleLabels name="label" toName="image">
+        <Label value="car" category="10"/>
+        <Label value="person" category="2"/>
+        <Label value="truck" category="7"/>
+      </RectangleLabels>
+    </View>
+    """.strip()
+    input_payload = [
+        {
+            "id": 1,
+            "data": {"image": "not-downloaded.jpg"},
+            "annotations": [
+                {
+                    "id": 1,
+                    "result": [
+                        {
+                            "id": "r1",
+                            "type": "rectanglelabels",
+                            "value": {
+                                "x": 10,
+                                "y": 10,
+                                "width": 20,
+                                "height": 20,
+                                "rotation": 0,
+                                "rectanglelabels": ["car"],
+                            },
+                            "to_name": "image",
+                            "from_name": "label",
+                            "original_width": 1000,
+                            "original_height": 500,
+                        },
+                        {
+                            "id": "r2",
+                            "type": "rectanglelabels",
+                            "value": {
+                                "x": 30,
+                                "y": 30,
+                                "width": 10,
+                                "height": 10,
+                                "rotation": 0,
+                                "rectanglelabels": ["truck"],
+                            },
+                            "to_name": "image",
+                            "from_name": "label",
+                            "original_width": 1000,
+                            "original_height": 500,
+                        },
+                    ],
+                }
+            ],
+        }
+    ]
+    input_path = temp_out_dir / "input.json"
+    input_path.write_text(json.dumps(input_payload))
+
+    converter = Converter(config=config, project_dir=PROJECT_DIR, download_resources=False)
+    converter.convert_to_coco(str(input_path), str(temp_out_dir), output_image_dir=str(temp_out_dir / "images"), is_dir=False)
+
+    coco_path = next(temp_out_dir.glob("*.json"))
+    coco = json.loads(coco_path.read_text())
+
+    category_id_by_name = {cat["name"]: cat["id"] for cat in coco["categories"]}
+    assert category_id_by_name["car"] == 10
+    assert category_id_by_name["person"] == 2
+    assert category_id_by_name["truck"] == 7
+    assert all(isinstance(cat["id"], int) for cat in coco["categories"])
+
+    annotation_category_ids = {ann["category_id"] for ann in coco["annotations"]}
+    assert annotation_category_ids == {10, 7}