Update app.py

app.py

@@ -1,26 +1,31 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import torch
-from ultralytics import YOLO
+import onnxruntime
 
-model = YOLO("best.onnx")
+# Download the ONNX model from Hugging Face
+model_path = hf_hub_download(repo_id="mkhug98/EchoYolo", filename="best.onnx")
+
+# Load the ONNX model
+session = onnxruntime.InferenceSession(model_path)
 
 # Function to perform object detection
 def detect_objects(image):
-    # Perform inference with the YOLOv5 model
-    results = model(image)
-    
-    # Get the detected bounding boxes and labels
-    bboxes = results.xyxy[0].cpu().numpy()
-    labels = results.names
+    # Preprocess the image
+    image = image.resize((640, 640))  # Resize the image to the expected input size
+    input_data = image.transpose(2, 0, 1).numpy()  # Rearrange the dimensions for ONNX
+
+    # Perform inference with the ONNX model
+    outputs = session.run(None, {"images": input_data.astype("float32")})
+    bboxes, scores, class_ids = outputs
 
     # Create a list of dictionaries for each detected object
     detections = []
-    for bbox, label in zip(bboxes, labels):
-        x1, y1, x2, y2, conf, cls = bbox
+    for bbox, score, class_id in zip(bboxes[0], scores[0], class_ids[0]):
+        x1, y1, x2, y2 = bbox
+        label = session.get_modelmeta().custom_metadata_map["names"][int(class_id)]
         detections.append({
             'label': label,
-            'confidence': float(conf),
+            'confidence': float(score),
             'x1': float(x1),
             'y1': float(y1),
             'x2': float(x2),