Spaces:
Running
Running
| import gradio as gr | |
| import PIL.Image as Image | |
| import tempfile | |
| import cv2 | |
| from ultralytics import ASSETS, YOLO | |
| # Load YOLOv8 model | |
| model = YOLO("yolov8n.pt") | |
| def predict_image(img, conf_threshold, iou_threshold): | |
| """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds.""" | |
| results = model.predict( | |
| source=img, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| show_labels=True, | |
| show_conf=True, | |
| imgsz=640, | |
| ) | |
| for r in results: | |
| im_array = r.plot() | |
| im = Image.fromarray(im_array[..., ::-1]) | |
| return im | |
| def predict_video(video_path, conf_threshold, iou_threshold): | |
| """Predicts objects in a video using a YOLOv8 model with adjustable confidence and IOU thresholds.""" | |
| # Create a temporary file to save the processed video | |
| temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) | |
| temp_output.close() | |
| # Load video | |
| cap = cv2.VideoCapture(video_path) | |
| # Get video properties | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| # Set up VideoWriter to save output video | |
| out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Perform inference on each frame | |
| results = model.predict( | |
| source=frame, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| show_labels=True, | |
| show_conf=True, | |
| imgsz=640, | |
| ) | |
| # Draw the results on the frame | |
| for r in results: | |
| frame = r.plot() | |
| # Write the frame to the output video | |
| out.write(frame) | |
| # Release resources | |
| cap.release() | |
| out.release() | |
| return temp_output.name | |
| def process_input(input_file, conf_threshold, iou_threshold, mode): | |
| """Handles both image and video inference based on the selected mode.""" | |
| if mode == "Image": | |
| return predict_image(input_file, conf_threshold, iou_threshold) | |
| elif mode == "Video": | |
| return predict_video(input_file.name, conf_threshold, iou_threshold) | |
| # Create Gradio interface | |
| iface = gr.Interface( | |
| fn=process_input, | |
| inputs=[ | |
| gr.File(label="Upload Image or Video File"), # Use a generic File input for both image and video | |
| gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"), | |
| gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"), | |
| gr.Radio(choices=["Image", "Video"], label="Select Mode", value="Image"), | |
| ], | |
| outputs=gr.Image(type="pil", label="Result") if gr.Radio else gr.Video(type="file", label="Result"), | |
| title="Ultralytics Gradio Application 🚀", | |
| description="Upload images or videos for inference. The Ultralytics YOLOv8n model is used by default.", | |
| examples=[ | |
| [ASSETS / "bus.jpg", 0.25, 0.45, "Image"], | |
| [ASSETS / "zidane.jpg", 0.25, 0.45, "Image"], | |
| ], | |
| ) | |
| iface.launch(share=True) | |