Use supervision library to visualize predictions

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,27 +46,8 @@ 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)
-    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
+    annotated_image = results[0].plot()
+    return annotated_image[:, :, ::-1]
 
 
 def app():

requirements.txt

@@ -6,5 +6,4 @@ gradio==4.42.0
 gradio_client==1.3.0
 gradio_image_prompter==0.1.0
 huggingface-hub==0.26.3
-fastapi==0.112.2
-supervision==0.25.1
+fastapi==0.112.2