with class remapping for YOLO-NAS

app.py

@@ -1,6 +1,5 @@
 from typing import Tuple
 
-import spaces
 import gradio as gr
 import numpy as np
 import supervision as sv
@@ -11,17 +10,17 @@ MARKDOWN = """
 
 Welcome to YOLO-Arena! This demo showcases the performance of various YOLO models:
 
-- YOLOv8
-- YOLOv9
-- YOLOv10
-- YOLO-NAS
+- YOLOv8 [[code](https://github.com/ultralytics/ultralytics)]
+- YOLOv9 [[code](https://github.com/WongKinYiu/yolov9)]
+- YOLOv10 [[code](https://github.com/THU-MIG/yolov10)]
+- YOLO-NAS [[code](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md)]
 
 Powered by Roboflow [Inference](https://github.com/roboflow/inference) and 
 [Supervision](https://github.com/roboflow/supervision).
 """
 
 IMAGE_EXAMPLES = [
-    ['https://media.roboflow.com/dog.jpeg', 0.3]
+    ['https://media.roboflow.com/dog.jpeg', 0.4]
 ]
 
 YOLO_V8_MODEL = get_model(model_id="yolov8m-640")
@@ -29,6 +28,17 @@ YOLO_NAS_MODEL = get_model(model_id="coco/15")
 YOLO_V9_MODEL = get_model(model_id="coco/17")
 YOLO_V10_MODEL = get_model(model_id="coco/22")
 
+YOLO_NAS_TO_COCO_CLASS_ID_MAPPING = {
+    49: 0, 9: 1, 18: 2, 44: 3, 0: 4, 16: 5, 73: 6, 74: 7, 11: 8, 72: 9, 31: 10, 63: 11,
+    48: 12, 8: 13, 10: 14, 20: 15, 28: 16, 37: 17, 56: 18, 25: 19, 30: 20, 6: 21,
+    79: 22, 34: 23, 2: 24, 76: 25, 36: 26, 68: 27, 64: 28, 33: 29, 59: 30, 60: 31,
+    62: 32, 40: 33, 4: 34, 5: 35, 58: 36, 65: 37, 67: 38, 13: 39, 78: 40, 26: 41,
+    32: 42, 41: 43, 61: 44, 14: 45, 3: 46, 1: 47, 54: 48, 46: 49, 15: 50, 19: 51,
+    38: 52, 50: 53, 29: 54, 17: 55, 22: 56, 24: 57, 51: 58, 7: 59, 27: 60, 70: 61,
+    75: 62, 42: 63, 45: 64, 53: 65, 39: 66, 21: 67, 43: 68, 47: 69, 69: 70, 57: 71,
+    52: 72, 12: 73, 23: 74, 77: 75, 55: 76, 66: 77, 35: 78, 71: 79
+}
+
 LABEL_ANNOTATORS = sv.LabelAnnotator(text_color=sv.Color.black())
 BOUNDING_BOX_ANNOTATORS = sv.BoundingBoxAnnotator()
 
@@ -37,7 +47,8 @@ def detect_and_annotate(
     model,
     input_image: np.ndarray,
     confidence_threshold: float,
-    iou_threshold: float
+    iou_threshold: float,
+    class_id_mapping: dict = None
 ) -> np.ndarray:
     result = model.infer(
         input_image,
@@ -45,15 +56,28 @@ def detect_and_annotate(
         iou_threshold=iou_threshold
     )[0]
     detections = sv.Detections.from_inference(result)
+
+    if class_id_mapping:
+        detections.class_id = np.array([
+            class_id_mapping[class_id]
+            for class_id
+            in detections.class_id
+        ])
+
+    labels = [
+        f"{class_name} ({confidence:.2f})"
+        for class_name, confidence
+        in zip(detections['class_name'], detections.confidence)
+    ]
+
     annotated_image = input_image.copy()
     annotated_image = BOUNDING_BOX_ANNOTATORS.annotate(
         scene=annotated_image, detections=detections)
     annotated_image = LABEL_ANNOTATORS.annotate(
-        scene=annotated_image, detections=detections)
+        scene=annotated_image, detections=detections, labels=labels)
     return annotated_image
 
 
-@spaces.GPU(duration=200)
 def process_image(
     input_image: np.ndarray,
     confidence_threshold: float,
@@ -62,7 +86,11 @@ def process_image(
     yolo_v8_annotated_image = detect_and_annotate(
         YOLO_V8_MODEL, input_image, confidence_threshold, iou_threshold)
     yolo_nas_annotated_image = detect_and_annotate(
-        YOLO_NAS_MODEL, input_image, confidence_threshold, iou_threshold)
+        YOLO_NAS_MODEL,
+        input_image,
+        confidence_threshold,
+        iou_threshold,
+        YOLO_NAS_TO_COCO_CLASS_ID_MAPPING)
     yolo_v9_annotated_image = detect_and_annotate(
         YOLO_V9_MODEL, input_image, confidence_threshold, iou_threshold)
     yolo_10_annotated_image = detect_and_annotate(