import cv2 import numpy as np import pandas as pd import plotly.express as px import plotly.graph_objects as go import torch import gradio as gr import os import time from scipy.optimize import curve_fit import sys # Add yolov5 directory to sys.path sys.path.append(os.path.join(os.path.dirname(__file__), "yolov5")) # Import YOLOv5 modules from models.experimental import attempt_load from utils.general import non_max_suppression, xywh2xyxy # Cricket pitch dimensions (in meters) PITCH_LENGTH = 20.12 # Length of cricket pitch (stumps to stumps) PITCH_WIDTH = 3.05 # Width of pitch STUMP_HEIGHT = 0.71 # Stump height STUMP_WIDTH = 0.2286 # Stump width (including bails) # Model input size (adjust if best.pt was trained with a different size) MODEL_INPUT_SIZE = (640, 640) # (height, width) FRAME_SKIP = 2 # Process every 2nd frame MIN_DETECTIONS = 10 # Stop after 10 detections BATCH_SIZE = 4 # Process 4 frames at a time # Load model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = attempt_load("best.pt") # Load without map_location model.to(device).eval() # Move model to device and set to evaluation mode # Function to process video and detect ball def process_video(video_path): cap = cv2.VideoCapture(video_path) frame_rate = 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)) positions = [] frame_numbers = [] bounce_frame = None bounce_point = None batch_frames = [] batch_frame_nums = [] frame_count = 0 start_time = time.time() while cap.isOpened(): frame_num = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) ret, frame = cap.read() if not ret: break # Skip frames if frame_count % FRAME_SKIP != 0: frame_count += 1 continue # Resize frame to model input size frame = cv2.resize(frame, MODEL_INPUT_SIZE, interpolation=cv2.INTER_AREA) batch_frames.append(frame) batch_frame_nums.append(frame_num) frame_count += 1 # Process batch when full or at end if len(batch_frames) == BATCH_SIZE or not ret: # Preprocess batch batch = [cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in batch_frames] batch = np.stack(batch) # [batch_size, H, W, 3] batch = torch.from_numpy(batch).to(device).float() / 255.0 batch = batch.permute(0, 3, 1, 2) # [batch_size, 3, H, W] # Run inference frame_start_time = time.time() with torch.no_grad(): pred = model(batch)[0] pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45) print(f"Batch inference time: {time.time() - frame_start_time:.2f}s for {len(batch_frames)} frames") # Process detections for i, det in enumerate(pred): if det is not None and len(det): det = xywh2xyxy(det) # Convert to [x1, y1, x2, y2] for *xyxy, conf, cls in det: x_center = (xyxy[0] + xyxy[2]) / 2 y_center = (xyxy[1] + xyxy[3]) / 2 # Scale coordinates back to original frame size x_center = x_center * frame_width / MODEL_INPUT_SIZE[1] y_center = y_center * frame_height / MODEL_INPUT_SIZE[0] positions.append((x_center.item(), y_center.item())) frame_numbers.append(batch_frame_nums[i]) # Detect bounce (lowest y_center point) if bounce_frame is None or y_center > positions[bounce_frame][1]: bounce_frame = len(frame_numbers) - 1 bounce_point = (x_center.item(), y_center.item()) batch_frames = [] batch_frame_nums = [] # Early termination if len(positions) >= MIN_DETECTIONS: break cap.release() print(f"Total video processing time: {time.time() - start_time:.2f}s") return positions, frame_numbers, bounce_point, frame_rate, frame_width, frame_height # Polynomial function for trajectory fitting def poly_func(x, a, b, c): return a * x**2 + b * x + c # Predict trajectory and LBW decision def predict_trajectory(positions, frame_numbers, frame_width, frame_height): if len(positions) < 3: return None, "Insufficient detections for trajectory prediction" x_coords = [p[0] for p in positions] y_coords = [p[1] for p in positions] frames = np.array(frame_numbers) # Fit polynomial to x and y coordinates try: popt_x, _ = curve_fit(poly_func, frames, x_coords) popt_y, _ = curve_fit(poly_func, frames, y_coords) except: return None, "Failed to fit trajectory" # Extrapolate to stumps frame_max = max(frames) + 10 future_frames = np.linspace(min(frames), frame_max, 100) x_pred = poly_func(future_frames, *popt_x) y_pred = poly_func(future_frames, *popt_y) # Check if trajectory hits stumps stump_x = frame_width / 2 stump_y = frame_height stump_hit = False for x, y in zip(x_pred, y_pred): if abs(y - stump_y) < 50 and abs(x - stump_x) < STUMP_WIDTH * frame_width / PITCH_WIDTH: stump_hit = True break lbw_decision = "OUT" if stump_hit else "NOT OUT" return list(zip(future_frames, x_pred, y_pred)), lbw_decision # Map pitch location def map_pitch(bounce_point, frame_width, frame_height): if bounce_point is None: return None, "No bounce detected" x, y = bounce_point pitch_x = (x / frame_width) * PITCH_WIDTH - PITCH_WIDTH / 2 pitch_y = (1 - y / frame_height) * PITCH_LENGTH return pitch_x, pitch_y # Estimate ball speed def estimate_speed(positions, frame_numbers, frame_rate, frame_width): if len(positions) < 2: return None, "Insufficient detections for speed estimation" distances = [] for i in range(1, len(positions)): x1, y1 = positions[i-1] x2, y2 = positions[i] pixel_dist = np.sqrt((x2 - x1)**2 + (y2 - y1)**2) distances.append(pixel_dist) pixel_to_meter = PITCH_LENGTH / frame_width distances_m = [d * pixel_to_meter for d in distances] time_interval = 1 / frame_rate speeds = [d / time_interval for d in distances_m] avg_speed_kmh = np.mean(speeds) * 3.6 return avg_speed_kmh, "Speed calculated successfully" # Create pitch map visualization def create_pitch_map(pitch_x, pitch_y): fig = go.Figure() fig.add_shape( type="rect", x0=-PITCH_WIDTH/2, y0=0, x1=PITCH_WIDTH/2, y1=PITCH_LENGTH, line=dict(color="Green"), fillcolor="Green", opacity=0.3 ) fig.add_shape( type="rect", x0=-STUMP_WIDTH/2, y0=PITCH_LENGTH-0.1, x1=STUMP_WIDTH/2, y1=PITCH_LENGTH, line=dict(color="Brown"), fillcolor="Brown" ) if pitch_x is not None and pitch_y is not None: fig.add_trace(go.Scatter(x=[pitch_x], y=[pitch_y], mode="markers", marker=dict(size=10, color="Red"), name="Bounce Point")) fig.update_layout( title="Pitch Map", xaxis_title="Width (m)", yaxis_title="Length (m)", xaxis_range=[-PITCH_WIDTH/2, PITCH_WIDTH/2], yaxis_range=[0, PITCH_LENGTH] ) return fig # Main Gradio function def drs_analysis(video): # Video is a file path (string) in Hugging Face Spaces video_path = video if isinstance(video, str) else "temp_video.mp4" if not isinstance(video, str): with open(video_path, "wb") as f: f.write(video.read()) positions, frame_numbers, bounce_point, frame_rate, frame_width, frame_height = process_video(video_path) if not positions: return None, None, "No ball detected in video", None trajectory, lbw_decision = predict_trajectory(positions, frame_numbers, frame_width, frame_height) if trajectory is None: return None, None, lbw_decision, None pitch_x, pitch_y = map_pitch(bounce_point, frame_width, frame_height) speed_kmh, speed_status = estimate_speed(positions, frame_numbers, frame_rate, frame_width) trajectory_df = pd.DataFrame(trajectory, columns=["Frame", "X", "Y"]) fig_traj = px.line(trajectory_df, x="X", y="Y", title="Ball Trajectory (Pixel Coordinates)") fig_traj.update_yaxes(autorange="reversed") fig_pitch = create_pitch_map(pitch_x, pitch_y) if not isinstance(video, str): os.remove(video_path) return fig_traj, fig_pitch, f"LBW Decision: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h", video_path # Gradio interface with gr.Blocks() as demo: gr.Markdown("## Cricket DRS Analysis") video_input = gr.Video(label="Upload Video Clip") btn = gr.Button("Analyze") trajectory_output = gr.Plot(label="Ball Trajectory") pitch_output = gr.Plot(label="Pitch Map") text_output = gr.Textbox(label="Analysis Results") video_output = gr.Video(label="Processed Video") btn.click(drs_analysis, inputs=video_input, outputs=[trajectory_output, pitch_output, text_output, video_output]) if __name__ == "__main__": demo.launch()