Spaces:
Running
Running
| import gradio as gr | |
| import cv2 | |
| from ultralytics import YOLO, solutions | |
| # Initialize the YOLO model | |
| model = YOLO("yolov8s.pt") | |
| def process_video(video_path, analytics_type): | |
| cap = cv2.VideoCapture(video_path) | |
| assert cap.isOpened(), "Error reading video file" | |
| # Get video properties | |
| w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) | |
| output_filename = f"{analytics_type}_output.avi" | |
| out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h)) | |
| # Set up analytics based on the selected type | |
| analytics = solutions.Analytics( | |
| type=analytics_type, | |
| writer=out, | |
| im0_shape=(w, h), | |
| view_img=False | |
| ) | |
| clswise_count = {} | |
| frame_count = 0 | |
| while cap.isOpened(): | |
| success, frame = cap.read() | |
| if success: | |
| frame_count += 1 | |
| results = model.track(frame, persist=True, verbose=False) | |
| if results[0].boxes.id is not None: | |
| boxes = results[0].boxes.xyxy.cpu() | |
| clss = results[0].boxes.cls.cpu().tolist() | |
| for box, cls in zip(boxes, clss): | |
| if model.names[int(cls)] in clswise_count: | |
| clswise_count[model.names[int(cls)]] += 1 | |
| else: | |
| clswise_count[model.names[int(cls)]] = 1 | |
| # Update analytics based on type | |
| if analytics_type == "line": | |
| total_counts = sum(clswise_count.values()) | |
| analytics.update_line(frame_count, total_counts) | |
| elif analytics_type == "multiple_line": | |
| analytics.update_multiple_lines(clswise_count, list(clswise_count.keys()), frame_count) | |
| elif analytics_type == "pie": | |
| analytics.update_pie(clswise_count) | |
| elif analytics_type == "area": | |
| analytics.update_area(frame_count, clswise_count) | |
| clswise_count = {} # Reset for next frame | |
| else: | |
| break | |
| cap.release() | |
| out.release() | |
| return output_filename # Return the output video file | |
| def gradio_app(video, analytics_type): | |
| # Save uploaded video locally | |
| video_path = video.name | |
| output_video = process_video(video_path, analytics_type) | |
| # Return processed video for display | |
| return output_video | |
| # Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# YOLO Video Processing App") | |
| with gr.Row(): | |
| video_input = gr.Video(label="Upload Video", type="file") | |
| analytics_dropdown = gr.Dropdown( | |
| ["line", "multiple_line", "pie", "area"], | |
| label="Select Analytics Type", | |
| value="line" | |
| ) | |
| output_video = gr.Video(label="Processed Output") | |
| # Button to start processing | |
| submit_btn = gr.Button("Process Video") | |
| # Define the output when the button is clicked | |
| submit_btn.click(gradio_app, inputs=[video_input, analytics_dropdown], outputs=output_video) | |
| # Launch the Gradio app | |
| demo.launch() |