Use supervision to visualize model results

app.py

@@ -1,6 +1,6 @@
-import torch
 import numpy as np
 import gradio as gr
+import supervision as sv
 from scipy.ndimage import binary_fill_holes
 from ultralytics import YOLOE
 from ultralytics.utils.torch_utils import smart_inference_mode
@@ -46,8 +46,27 @@ def yoloe_inference(image, prompts, target_image, model_id, image_size, conf_thr
         model.model.model[-1].max_det = 1000
 
     results = model.predict(source=image, imgsz=image_size, conf=conf_thresh, iou=iou_thresh, **kwargs)
-    annotated_image = results[0].plot()
-    return annotated_image[:, :, ::-1]
+    detections = sv.Detections.from_ultralytics(results[0])
+
+    resolution_wh = image.size
+    thickness = sv.calculate_optimal_line_thickness(resolution_wh=resolution_wh)
+    text_scale = sv.calculate_optimal_text_scale(resolution_wh=resolution_wh)
+
+    labels = [
+        f"{class_name} {confidence:.2f}"
+        for class_name, confidence
+        in zip(detections['class_name'], detections.confidence)
+    ]
+
+    annotated_image = image.copy()
+    annotated_image = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX, opacity=0.4).annotate(
+        scene=annotated_image, detections=detections)
+    annotated_image = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX, thickness=thickness).annotate(
+        scene=annotated_image, detections=detections)
+    annotated_image = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX, text_scale=text_scale, smart_position=True).annotate(
+        scene=annotated_image, detections=detections, labels=labels)
+
+    return annotated_image
 
 
 def app():

requirements.txt

@@ -1,9 +1,10 @@
 git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/CLIP
 git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/ml-mobileclip
-git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/lvis-api
+# git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/lvis-api
 git+https://github.com/THU-MIG/yoloe.git
 gradio==4.42.0
 gradio_client==1.3.0
 gradio_image_prompter==0.1.0
 huggingface-hub==0.26.3
-fastapi==0.112.2
+fastapi==0.112.2
+supervision==0.25.1