Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| from PIL import Image | |
| import gradio as gr | |
| from huggingface_hub import snapshot_download | |
| import os | |
| import cv2 | |
| import numpy as np | |
| def load_model(repo_id): | |
| """Download and load the YOLO model.""" | |
| download_dir = snapshot_download(repo_id) | |
| path = os.path.join(download_dir, "model_- 5 january 2025 0_48_openvino_model") | |
| detection_model = YOLO(path, task='detect') | |
| return detection_model | |
| def predict_image(pilimg, conf_threshold, iou_threshold): | |
| """Process an image with user-defined thresholds.""" | |
| result = detection_model.predict(pilimg, conf=conf_threshold, iou=iou_threshold) | |
| img_bgr = result[0].plot() | |
| out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # Convert to RGB PIL image | |
| return out_pilimg | |
| def predict_video(video_file, conf_threshold, iou_threshold, start_time, end_time): | |
| """Process a video with user-defined thresholds and time range.""" | |
| cap = cv2.VideoCapture(video_file) | |
| # Get video properties | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| output_path = "output_video.mp4" | |
| # Calculate frame range based on start and end times | |
| start_frame = int(start_time * fps) if start_time else 0 | |
| end_frame = int(end_time * fps) if end_time else total_frames | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) # Start processing at the desired frame | |
| # Initialize VideoWriter | |
| out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height)) | |
| while cap.isOpened(): | |
| current_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) | |
| if current_frame >= end_frame: | |
| break | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Resize frame to 640x640 for faster inference | |
| resized_frame = cv2.resize(frame, (640, 640)) | |
| # Perform detection on resized frame | |
| result = detection_model.predict(resized_frame, conf=conf_threshold, iou=iou_threshold) | |
| # Get processed frame and resize back to original dimensions | |
| output_frame = result[0].plot() | |
| output_frame = cv2.resize(output_frame, (frame_width, frame_height)) | |
| # Write the processed frame to output video | |
| out.write(output_frame) | |
| cap.release() | |
| return output_path | |
| # Load the model | |
| REPO_ID = "WgJL/107_Assignment.v5i.yolov11" | |
| detection_model = load_model(REPO_ID) | |
| # Gradio UI | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## Lemons and Crabs Detection") | |
| with gr.Tab("Image Input"): | |
| img_input = gr.Image(type="pil", label="Upload an Image") | |
| conf_slider_img = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold") | |
| iou_slider_img = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="IoU Threshold") | |
| img_output = gr.Image(type="pil", label="Processed Image") | |
| img_submit = gr.Button("Process Image") | |
| img_submit.click( | |
| predict_image, | |
| inputs=[img_input, conf_slider_img, iou_slider_img], | |
| outputs=img_output | |
| ) | |
| with gr.Tab("Video Input"): | |
| video_input = gr.Video(label="Upload a Video") | |
| conf_slider_video = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold") | |
| iou_slider_video = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="IoU Threshold") | |
| start_time = gr.Number(value=0, label="Start Time (seconds)") | |
| end_time = gr.Number(value=0, label="End Time (seconds, 0 for full video)") | |
| video_output = gr.Video(label="Processed Video") | |
| video_submit = gr.Button("Process Video") | |
| video_submit.click( | |
| predict_video, | |
| inputs=[video_input, conf_slider_video, iou_slider_video, start_time, end_time], | |
| outputs=video_output | |
| ) | |
| demo.launch(share=True) | |