Update app.py

app.py

@@ -1,154 +1,287 @@
 import cv2
 import numpy as np
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
 import torch
-from ultralytics import YOLO
 import gradio as gr
-from scipy.interpolate import interp1d
-import uuid
 import os
+import time
+from scipy.optimize import curve_fit
+import sys
 
-# Load the trained YOLOv8n model from the Space's root directory
-model = YOLO("best.pt")  # Assumes best.pt is in the same directory as app.py
+# Add yolov5 directory to sys.path
+sys.path.append(os.path.join(os.path.dirname(__file__), "yolov5"))
 
-# Constants for LBW decision and video processing
-STUMPS_WIDTH = 0.2286  # meters (width of stumps)
-BALL_DIAMETER = 0.073  # meters (approx. cricket ball diameter)
-FRAME_RATE = 30  # Input video frame rate
-SLOW_MOTION_FACTOR = 6  # For very slow motion (6x slower)
-CONF_THRESHOLD = 0.3  # Lowered confidence threshold for better detection
+# 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 yolov5s.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
+SLOW_MOTION_FACTOR = 3  # Duplicate each frame 3 times for slow motion
+
+# Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = attempt_load("best.pt")  # Load yolov5s.pt
+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):
-    # Initialize video capture
-    if not os.path.exists(video_path):
-        return [], [], "Error: Video file not found"
     cap = cv2.VideoCapture(video_path)
-    frames = []
-    ball_positions = []
-    debug_log = []
-
+    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
-        frames.append(frame.copy())  # Store original frame
-        # Detect ball using the trained YOLOv8n model
-        results = model.predict(frame, conf=CONF_THRESHOLD)
-        detections = 0
-        for detection in results[0].boxes:
-            if detection.cls == 0:  # Assuming class 0 is the ball
-                detections += 1
-                x1, y1, x2, y2 = detection.xyxy[0].cpu().numpy()
-                ball_positions.append([(x1 + x2) / 2, (y1 + y2) / 2])
-                # Draw bounding box on frame for visualization
-                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
-        frames[-1] = frame  # Update frame with bounding box
-        debug_log.append(f"Frame {frame_count}: {detections} ball detections")
+
+        # 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
 
-    if not ball_positions:
-        debug_log.append("No balls detected in any frame")
-    else:
-        debug_log.append(f"Total ball detections: {len(ball_positions)}")
+# Polynomial function for trajectory fitting
+def poly_func(x, a, b, c):
+    return a * x**2 + b * x + c
 
-    return frames, ball_positions, "\n".join(debug_log)
+# Predict trajectory and wicket inline path
+def predict_trajectory(positions, frame_numbers, frame_width, frame_height):
+    if len(positions) < 3:
+        return None, None, "Insufficient detections for trajectory prediction"
 
-def estimate_trajectory(ball_positions, frames):
-    # Simplified physics-based trajectory projection
-    if len(ball_positions) < 2:
-        return None, None, "Error: Fewer than 2 ball detections for trajectory"
-    # Extract x, y coordinates
-    x_coords = [pos[0] for pos in ball_positions]
-    y_coords = [pos[1] for pos in ball_positions]
-    times = np.arange(len(ball_positions)) / FRAME_RATE
+    x_coords = [p[0] for p in positions]
+    y_coords = [p[1] for p in positions]
+    frames = np.array(frame_numbers)
 
-    # Interpolate to smooth trajectory
+    # Fit polynomial to x and y coordinates
     try:
-        fx = interp1d(times, x_coords, kind='linear', fill_value="extrapolate")
-        fy = interp1d(times, y_coords, kind='quadratic', fill_value="extrapolate")
-    except Exception as e:
-        return None, None, f"Error in trajectory interpolation: {str(e)}"
-
-    # Project trajectory forward (0.5 seconds post-impact)
-    t_future = np.linspace(times[-1], times[-1] + 0.5, 10)
-    x_future = fx(t_future)
-    y_future = fy(t_future)
-
-    return list(zip(x_future, y_future)), t_future, "Trajectory estimated successfully"
-
-def lbw_decision(ball_positions, trajectory, frames):
-    # Simplified LBW logic
-    if not frames:
-        return "Error: No frames processed", None
-    if not trajectory or len(ball_positions) < 2:
-        return "Not enough data (insufficient ball detections)", None
-
-    # Assume stumps are at the bottom center of the frame (calibration needed)
-    frame_height, frame_width = frames[0].shape[:2]
-    stumps_x = frame_width / 2
-    stumps_y = frame_height * 0.9  # Approximate stumps position
-    stumps_width_pixels = frame_width * (STUMPS_WIDTH / 3.0)  # Assume 3m pitch width
-
-    # Check pitching point (first detected position)
-    pitch_x, pitch_y = ball_positions[0]
-    if pitch_x < stumps_x - stumps_width_pixels / 2 or pitch_x > stumps_x + stumps_width_pixels / 2:
-        return "Not Out (Pitched outside line)", None
-
-    # Check impact point (last detected position)
-    impact_x, impact_y = ball_positions[-1]
-    if impact_x < stumps_x - stumps_width_pixels / 2 or impact_x > stumps_x + stumps_width_pixels / 2:
-        return "Not Out (Impact outside line)", None
-
-    # Check trajectory hitting stumps
-    for x, y in trajectory:
-        if abs(x - stumps_x) < stumps_width_pixels / 2 and abs(y - stumps_y) < frame_height * 0.1:
-            return "Out", trajectory
-    return "Not Out (Missing stumps)", trajectory
-
-def generate_slow_motion(frames, trajectory, output_path):
-    # Generate very slow-motion video with ball detection and trajectory overlay
-    if not frames:
-        return None
+        popt_x, _ = curve_fit(poly_func, frames, x_coords)
+        popt_y, _ = curve_fit(poly_func, frames, y_coords)
+    except:
+        return None, 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)
+
+    # Wicket inline path (center line toward stumps)
+    stump_x = frame_width / 2
+    stump_y = frame_height
+    inline_x = np.linspace(min(x_coords), stump_x, 100)
+    inline_y = np.interp(inline_x, x_pred, y_pred)
+
+    # Check if trajectory hits stumps
+    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)), list(zip(inline_x, inline_y)), 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"
+
+# Main Gradio function with video overlay and slow motion
+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())
+
+    # Process video for detections
+    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
+
+    # Predict trajectory and wicket path
+    trajectory, inline_path, 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)
+
+    # Create output video with overlays and slow motion
+    output_path = "output_video.mp4"
+    cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frames[0].shape[1], frames[0].shape[0]))
+    out = cv2.VideoWriter(output_path, fourcc, frame_rate, (frame_width, frame_height))
+
+    frame_count = 0
+    positions_dict = dict(zip(frame_numbers, positions))
 
-    for frame in frames:
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Skip frames for consistency with detection
+        if frame_count % FRAME_SKIP != 0:
+            frame_count += 1
+            continue
+
+        # Overlay ball trajectory (red) and wicket inline path (blue)
+        if frame_count in positions_dict:
+            cv2.circle(frame, (int(positions_dict[frame_count][0]), int(positions_dict[frame_count][1])), 5, (0, 0, 255), -1)  # Red dot
         if trajectory:
-            for x, y in trajectory:
-                cv2.circle(frame, (int(x), int(y)), 5, (255, 0, 0), -1)  # Blue dots for trajectory
-        for _ in range(SLOW_MOTION_FACTOR):  # Duplicate frames for very slow motion
+            traj_x = [int(t[1]) for t in trajectory if t[0] >= frame_count]
+            traj_y = [int(t[2]) for t in trajectory if t[0] >= frame_count]
+            if traj_x and traj_y:
+                for i in range(1, len(traj_x)):
+                    cv2.line(frame, (traj_x[i-1], traj_y[i-1]), (traj_x[i], traj_y[i]), (0, 0, 255), 2)  # Red line
+        if inline_path:
+            inline_x = [int(x) for x, _ in inline_path]
+            inline_y = [int(y) for _, y in inline_path]
+            if inline_x and inline_y:
+                for i in range(1, len(inline_x)):
+                    cv2.line(frame, (inline_x[i-1], inline_y[i-1]), (inline_x[i], inline_y[i]), (255, 0, 0), 2)  # Blue line
+
+        # Overlay pitch map in top-right corner
+        if pitch_x is not None and pitch_y is not None:
+            map_width = 200
+            # Cap map_height to 25% of frame height to ensure it fits
+            map_height = min(int(map_width * PITCH_LENGTH / PITCH_WIDTH), frame_height // 4)
+            pitch_map = np.zeros((map_height, map_width, 3), dtype=np.uint8)
+            pitch_map[:] = (0, 255, 0)  # Green pitch
+            cv2.rectangle(pitch_map, (0, map_height-10), (map_width, map_height), (0, 51, 51), -1)  # Brown stumps
+            bounce_x = int((pitch_x + PITCH_WIDTH/2) / PITCH_WIDTH * map_width)
+            bounce_y = int((1 - pitch_y / PITCH_LENGTH) * map_height)
+            cv2.circle(pitch_map, (bounce_x, bounce_y), 5, (0, 0, 255), -1)  # Red bounce point
+            # Ensure overlay fits within frame
+            overlay_region = frame[0:map_height, frame_width-map_width:frame_width]
+            if overlay_region.shape[0] >= map_height and overlay_region.shape[1] >= map_width:
+                frame[0:map_height, frame_width-map_width:frame_width] = cv2.resize(pitch_map, (map_width, map_height))
+
+        # Add text annotations
+        text = f"LBW: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h"
+        cv2.putText(frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
+
+        # Write frame multiple times for slow motion
+        for _ in range(SLOW_MOTION_FACTOR):
             out.write(frame)
-    out.release()
-    return output_path
 
-def drs_review(video):
-    # Process video and generate DRS output
-    frames, ball_positions, debug_log = process_video(video)
-    if not frames:
-        return f"Error: Failed to process video\nDebug Log:\n{debug_log}", None
-    trajectory, _, trajectory_log = estimate_trajectory(ball_positions, frames)
-    decision, trajectory = lbw_decision(ball_positions, trajectory, frames)
+        frame_count += 1
+
+    cap.release()
+    out.release()
 
-    # Generate slow-motion replay even if Trajectory fails
-    output_path = f"output_{uuid.uuid4()}.mp4"
-    slow_motion_path = generate_slow_motion(frames, trajectory, output_path)
+    if not isinstance(video, str):
+        os.remove(video_path)
 
-    # Combine debug logs for output
-    debug_output = f"{debug_log}\n{trajectory_log}"
-    return f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", slow_motion_path
+    return None, None, None, output_path
 
 # Gradio interface
-iface = gr.Interface(
-    fn=drs_review,
-    inputs=gr.Video(label="Upload Video Clip"),
-    outputs=[
-        gr.Textbox(label="DRS Decision and Debug Log"),
-        gr.Video(label="Very Slow-Motion Replay with Ball Detection and Trajectory")
-    ],
-    title="AI-Powered DRS for LBW in Local Cricket",
-    description="Upload a video clip of a cricket delivery to get an LBW decision and very slow-motion replay showing ball detection (green boxes) and trajectory (blue dots)."
-)
+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__":
-    iface.launch()
+    demo.launch()