Spaces:
Sleeping
Sleeping
| 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() |