Update app.py

app.py

@@ -1,32 +1,57 @@
 import gradio as gr
 import cv2
-from ultralytics import YOLO
+import numpy as np
+import onnxruntime as ort
 
-# Load the pretrained YOLOv5 model
-yolo_model = YOLO("yolov5su.pt")  # Load the small version of YOLOv5
-webcamera = cv2.VideoCapture(0)
+# Load the ONNX model using onnxruntime
+onnx_model_path = "aje6/ASL-Fingerspelling-Detection/onnx/Model_IV.onnx"  # Update with your ONNX model path
+session = ort.InferenceSession(onnx_model_path)
 
-# Function to perform object detection with a configurable confidence threshold
+# Function to perform object detection with the ONNX model
 def detect_objects(frame, confidence_threshold=0.5):
     # Convert the frame from BGR (OpenCV) to RGB
     image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
     
-    # Perform inference with YOLO, passing the confidence threshold
-    results = yolo_model(image, conf=confidence_threshold)  # Set confidence threshold
+    # Preprocessing: Resize and normalize the image
+    # Assuming YOLO model input is 640x640, update according to your model's input size
+    input_size = (640, 640)
+    image_resized = cv2.resize(image, input_size)
+    image_normalized = image_resized / 255.0  # Normalize to [0, 1]
+    image_input = np.transpose(image_normalized, (2, 0, 1))  # Change to CHW format
+    image_input = np.expand_dims(image_input, axis=0).astype(np.float32)  # Add batch dimension
+
+    # Perform inference
+    inputs = {session.get_inputs()[0].name: image_input}
+    outputs = session.run(None, inputs)
+    
+    # Assuming YOLO model outputs are in the form of [boxes, confidences, class_probs]
+    boxes, confidences, class_probs = outputs
+
+    # Post-processing: Filter boxes by confidence threshold
+    detections = []
+    for i, confidence in enumerate(confidences[0]):
+        if confidence >= confidence_threshold:
+            x1, y1, x2, y2 = boxes[0][i]
+            class_id = np.argmax(class_probs[0][i])  # Get class with highest probability
+            detections.append((x1, y1, x2, y2, confidence, class_id))
     
     # Draw bounding boxes and labels on the image
-    annotated_image = results.plot()  # This automatically draws boxes and labels
+    for (x1, y1, x2, y2, confidence, class_id) in detections:
+        color = (0, 255, 0)  # Green color for bounding boxes
+        cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
+        label = f"Class {class_id}: {confidence:.2f}"
+        cv2.putText(image, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     
     # Convert the image back to BGR for displaying in Gradio
-    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)
+    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     
-    return annotated_image
+    return image_bgr
 
 # Gradio interface to use the webcam for real-time object detection
 # Added a slider for the confidence threshold
 iface = gr.Interface(fn=detect_objects, 
                      inputs=[
-                         gr.Video(source=webcamera, type="numpy"),  # Webcam input
+                         gr.Video(source="webcam", type="numpy"),  # Webcam input
                          gr.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Confidence Threshold")  # Confidence slider
                      ],  
                      outputs="image")  # Show output image with bounding boxes