Spaces:
Sleeping
Sleeping
File size: 2,748 Bytes
2ee1f45 fd83372 2ee1f45 fd83372 3fed143 fd83372 2ee1f45 fd83372 2ee1f45 fd83372 9303397 fd83372 9303397 2ee1f45 9303397 2ee1f45 9303397 2ee1f45 9303397 2ee1f45 616a781 5ffc0a0 9a2da70 616a781 2ee1f45 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import gradio as gr
import cv2
import numpy as np
import onnxruntime as ort
# Load the ONNX model using onnxruntime
onnx_model_path = "Model_IV.onnx" # Update with your ONNX model path
session = ort.InferenceSession(onnx_model_path)
# 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)
# 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
# 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
# image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
return outputs
# 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(sources="webcam", type="numpy"), # Webcam input
inputs = gr.Image(sources=["webcam"], type="numpy"),
# gr.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Confidence Threshold") # Confidence slider
# ],
outputs="image") # Show output image with bounding boxes
iface.launch() |