refactor init_model function to accept model_id parameter and update predict endpoint to use dynamic model initialization; added supervision library to requirements

app.py

@@ -7,7 +7,7 @@ import os
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 import requests
-
+import supervision as sv
 
 ###
 
@@ -19,9 +19,8 @@ import requests
 #wget -q https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_blt.pt
 
 
-def init_model():
+def init_model(model_id: str):
     is_pf=True
-    model_id = "yoloe-11s"
     # Create a models directory if it doesn't exist
     os.makedirs("models", exist_ok=True)
     filename = f"{model_id}-seg.pt" if not is_pf else f"{model_id}-seg-pf.pt"
@@ -34,11 +33,16 @@ def init_model():
 
 app = FastAPI()
 
-# Initialize model at startup
-model = init_model()
 
 @app.post("/predict")
-async def predict(image_url: str, texts: str = "hat"):
+async def predict(image_url: str,
+                  texts: str = "hat",
+                  model_id: str = "yoloe-11l",
+                  conf: float = 0.25,
+                  iou: float = 0.7
+                  ):
+    # Initialize model at startup
+    model = init_model(model_id)
     # Set classes to filter
     class_list = [text.strip() for text in texts.split(',')]
     
@@ -51,7 +55,7 @@ async def predict(image_url: str, texts: str = "hat"):
     model.set_classes(class_list, text_embeddings)
     
     # Run inference with the PIL Image
-    results = model.predict(source=image, conf=0.25, iou=0.7)
+    results = model.predict(source=image, conf=conf, iou=iou)
     
     # Extract detection results
     result = results[0]
@@ -66,6 +70,25 @@ async def predict(image_url: str, texts: str = "hat"):
         }
         detections.append(detection)
     print(detections)
+    # 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 {"detections": detections}
 
 if __name__ == "__main__":

requirements.txt

@@ -1,5 +1,6 @@
 fastapi
 uvicorn[standard]
+supervision
 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