import gradio as gr import cv2 import numpy as np from ultralytics import solutions from tqdm import tqdm import tempfile import os # Global variables to store points and first frame points = [] first_frame = None video_path = None def extract_first_frame(video): """Extract first frame from uploaded video""" global first_frame, video_path, points points = [] # Reset points when new video is uploaded if video is None: return None, "Please upload a video first" video_path = video cap = cv2.VideoCapture(video) if not cap.isOpened(): return None, "Error: Could not open video file" ret, frame = cap.read() cap.release() if not ret: return None, "Error: Could not read video frame" first_frame = frame.copy() # Convert BGR to RGB for display frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) return frame_rgb, "Video loaded! Click on the image to select 2 points for the counting line." def select_points(evt: gr.SelectData, image): """Handle point selection on the image""" global points, first_frame if first_frame is None: return image, "Please upload a video first" if len(points) >= 2: return image, "Two points already selected. Click 'Reset Points' to start over." # Get click coordinates x, y = evt.index[0], evt.index[1] points.append((x, y)) # Draw points and line on image frame_display = first_frame.copy() # Draw all points for i, point in enumerate(points): cv2.circle(frame_display, point, 8, (0, 255, 0), -1) label = "Start" if i == 0 else "End" cv2.putText(frame_display, label, (point[0] + 10, point[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # Draw line if both points selected if len(points) == 2: cv2.line(frame_display, points[0], points[1], (0, 255, 0), 3) # Convert BGR to RGB for display frame_rgb = cv2.cvtColor(frame_display, cv2.COLOR_BGR2RGB) status = f"Points selected: {len(points)}/2" if len(points) == 2: status += " - Ready to process! Click 'Start Processing'" return frame_rgb, status def reset_points(): """Reset selected points""" global points, first_frame points = [] if first_frame is None: return None, "No video loaded" # Show original frame without points frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB) return frame_rgb, "Points reset. Click on the image to select new points." def process_video(progress=gr.Progress()): """Process video with selected counting line""" global points, video_path if video_path is None: return None, "Please upload a video first" if len(points) != 2: return None, f"Please select exactly 2 points. Currently selected: {len(points)}" # Create region points in the format expected by ObjectCounter region_points = [points[0], points[1]] # Open video cap = cv2.VideoCapture(video_path) if not cap.isOpened(): return None, "Error: Could not open video file" # Get video properties w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = int(cap.get(cv2.CAP_PROP_FPS)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Create temporary output file output_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name # Video writer fourcc = cv2.VideoWriter_fourcc(*'mp4v') video_writer = cv2.VideoWriter(output_path, fourcc, fps, (w, h)) # Initialize ObjectCounter counter = solutions.ObjectCounter( show=False, # Don't display during processing region=region_points, model="yolo11n.pt", classes=[2, 3], # car and motorcycle show_in=True, line_width=2, ) # Process video with progress bar frame_count = 0 progress(0, desc="Starting processing...") for frame_count in tqdm(range(total_frames), desc="Processing video"): ret, frame = cap.read() if not ret: break # Process frame results = counter(frame) # Write processed frame video_writer.write(results.plot_im) # Update progress progress((frame_count + 1) / total_frames, desc=f"Processing frame {frame_count + 1}/{total_frames}") # Release resources cap.release() video_writer.release() # Get final counts status = f"✅ Processing complete!\n" status += f"Total frames processed: {frame_count + 1}\n" status += f"Output saved to: {output_path}" return output_path, status # Create Gradio interface with gr.Blocks(title="Vehicle Counter", theme=gr.themes.Soft()) as app: gr.Markdown( """

🚗 Vehicle Counter

Upload a video, select a counting line by clicking two points, and process the video to count vehicles. """ ) with gr.Row(): with gr.Column(scale=1): video_input = gr.Video(label="Upload Video") load_btn = gr.Button("Load Video", variant="primary") gr.Markdown("### Instructions") gr.Markdown( """ 1. Upload a video file 2. Click 'Load Video' to extract first frame 3. Click on the image to select **Start** point 4. Click again to select **End** point 5. Click 'Start Processing' to count vehicles """ ) reset_btn = gr.Button("Reset Points", variant="secondary") process_btn = gr.Button("Start Processing", variant="primary", size="lg") with gr.Column(scale=2): frame_display = gr.Image(label="Video Frame - Click to select points", type="numpy") status_text = gr.Textbox(label="Status", lines=3, interactive=False) gr.Markdown("---") with gr.Row(): output_video = gr.Video(label="Processed Video with Counts") # Event handlers load_btn.click( fn=extract_first_frame, inputs=[video_input], outputs=[frame_display, status_text] ) frame_display.select( fn=select_points, inputs=[frame_display], outputs=[frame_display, status_text] ) reset_btn.click( fn=reset_points, inputs=[], outputs=[frame_display, status_text] ) process_btn.click( fn=process_video, inputs=[], outputs=[output_video, status_text] ) # Launch the app if __name__ == "__main__": app.launch()