Update app.py

app.py

@@ -34,13 +34,24 @@ def yolo_inference(image, model_id, conf_threshold, iou_threshold, max_detection
     model_path = download_models(model_id)
     model = YOLO(model_path)
     results = model(source=image, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
+    
+    # Get the detections and convert them to the supervision Detections format
     detections = sv.Detections.from_ultralytics(results)
     
+    # Prepare the labels
     labels = [
         f"{category_dict[class_id]} {confidence:.2f}"
         for class_id, confidence in zip(detections.class_id, detections.confidence)
     ]
-    annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)
+    
+
+    annotated_image = image.copy()  
+    draw = ImageDraw.Draw(annotated_image)
+
+    for label, (x1, y1, x2, y2) in zip(labels, detections.xyxy):  
+        color = random_color() 
+        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
+        draw.text((x1, y1), label, fill=color)
 
     return annotated_image