From d0b956918bfb5326d1f04f518b610c1a044001cc Mon Sep 17 00:00:00 2001
From: Corentin Vannier <corentin.vannier@cea.fr>
Date: Mon, 22 Jan 2024 15:18:53 +0100
Subject: [PATCH] feat(back): Replace category ID and name by single category
 field

---
 pixano_inference/pytorch/deeplabv3.py       | 3 +--
 pixano_inference/pytorch/maskrcnnv2.py      | 3 +--
 pixano_inference/pytorch/yolov5.py          | 5 +----
 pixano_inference/tensorflow/efficientdet.py | 3 +--
 pixano_inference/tensorflow/fasterrcnn.py   | 3 +--
 5 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/pixano_inference/pytorch/deeplabv3.py b/pixano_inference/pytorch/deeplabv3.py
index 050e126..b0bdef1 100644
--- 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
index ed69b7e..810d0cc 100644
--- 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
index 8e6df2f..5347b6d 100644
--- 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
index 3cd8f24..a963245 100644
--- 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
index b9ab5e6..b892baf 100644
--- 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]
                             ),
                         }