Update coco_ids_80to91() name

pixano · Jun 5, 2023 · 959203d · 959203d
1 parent 6891162
commit 959203d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pixano_inference/pytorch_models/yolov5.py b/pixano_inference/pytorch_models/yolov5.py
@@ -19,7 +19,7 @@
 from PIL import Image
 from pixano.core import arrow_types
 from pixano.inference import InferenceModel
-from pixano.transforms import coco_80to91, coco_names_91, normalize, xyxy_to_xywh
+from pixano.transforms import coco_ids_80to91, coco_names_91, normalize, xyxy_to_xywh
 
 
 class YOLOv5(InferenceModel):
@@ -95,8 +95,8 @@ def inference_batch(
                         bbox=normalize(xyxy_to_xywh(pred[0:4]), h, w),
                         bbox_confidence=float(pred[4]),
                         bbox_source=self.id,
-                        category_id=coco_80to91(pred[5] + 1),
-                        category_name=coco_names_91(coco_80to91(pred[5] + 1)),
+                        category_id=coco_ids_80to91(pred[5] + 1),
+                        category_name=coco_names_91(coco_ids_80to91(pred[5] + 1)),
                     )
                     for pred in img_output
                     if pred[4] > threshold