feat(back): Replace category ID and name by single category field

pixano · Jan 22, 2024 · d0b9569 · d0b9569
1 parent 14e558b
commit d0b9569
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 12 deletions.
diff --git a/pixano_inference/pytorch/deeplabv3.py b/pixano_inference/pytorch/deeplabv3.py
@@ -139,8 +139,7 @@ def preannotate(
                             "mask": CompressedRLE.from_mask(
                                 unmold_mask(mask)
                             ).to_dict(),
-                            "category_id": int(label),
-                            "category_name": voc_names(label),
+                            "category": voc_names(label),
                         }
                         for label, mask in zip(labels, masks)
                     ]

diff --git a/pixano_inference/pytorch/maskrcnnv2.py b/pixano_inference/pytorch/maskrcnnv2.py
@@ -140,8 +140,7 @@ def preannotate(
                             "mask": CompressedRLE.from_mask(
                                 unmold_mask(output["masks"][i])
                             ).to_dict(),
-                            "category_id": int(output["labels"][i]),
-                            "category_name": coco_names_91(output["labels"][i]),
+                            "category": coco_names_91(output["labels"][i]),
                         }
                         for i in range(len(output["scores"]))
                         if output["scores"][i] > threshold

diff --git a/pixano_inference/pytorch/yolov5.py b/pixano_inference/pytorch/yolov5.py
@@ -108,10 +108,7 @@ def preannotate(
                             )
                             .normalize(h, w)
                             .to_dict(),
-                            "category_id": coco_ids_80to91(pred[5] + 1),
-                            "category_name": coco_names_91(
-                                coco_ids_80to91(pred[5] + 1)
-                            ),
+                            "category": coco_names_91(coco_ids_80to91(pred[5] + 1)),
                         }
                         for pred in img_output
                         if pred[4] > threshold

diff --git a/pixano_inference/tensorflow/efficientdet.py b/pixano_inference/tensorflow/efficientdet.py
@@ -104,8 +104,7 @@ def preannotate(
                                 ],
                                 confidence=float(output["detection_scores"][0][i]),
                             ).to_dict(),
-                            "category_id": int(output["detection_classes"][0][i]),
-                            "category_name": coco_names_91(
+                            "category": coco_names_91(
                                 output["detection_classes"][0][i]
                             ),
                         }

diff --git a/pixano_inference/tensorflow/fasterrcnn.py b/pixano_inference/tensorflow/fasterrcnn.py
@@ -106,8 +106,7 @@ def preannotate(
                                 ],
                                 confidence=float(output["detection_scores"][0][i]),
                             ).to_dict(),
-                            "category_id": int(output["detection_classes"][0][i]),
-                            "category_name": coco_names_91(
+                            "category": coco_names_91(
                                 output["detection_classes"][0][i]
                             ),
                         }