Update app.py

app.py

@@ -3,28 +3,39 @@ import numpy as np
 import torch
 from ultralytics import YOLO
 import gradio as gr
-from scipy.interpolate import interp1d, CubicSpline
+from scipy.interpolate import interp1d
+import plotly.graph_objects as go
 import uuid
 import os
 
-# Load the trained YOLOv8n model
+# Load the trained YOLOv8n model with optimizations
 model = YOLO("best.pt")
+model.to('cuda' if torch.cuda.is_available() else 'cpu')  # Use GPU if available
 
 # 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 = 20  # Input video frame rate (reduced to 20 FPS)
-SLOW_MOTION_FACTOR = 3  # Adjusted for 20 FPS (slower playback without being too slow)
-CONF_THRESHOLD = 0.25  # Confidence threshold for detection
-IMPACT_ZONE_Y = 0.85  # Fraction of frame height where impact is likely (near stumps)
+FRAME_RATE = 20  # Input video frame rate
+SLOW_MOTION_FACTOR = 3  # For very slow motion (3x slower)
+CONF_THRESHOLD = 0.2  # Confidence threshold
+IMPACT_ZONE_Y = 0.85  # Fraction of frame height where impact is likely
+IMPACT_DELTA_Y = 50  # Pixels for detecting sudden y-position change
+PITCH_LENGTH = 20.12  # meters (standard cricket pitch length)
+STUMPS_HEIGHT = 0.71  # meters (stumps height)
+CAMERA_HEIGHT = 2.0  # meters (assumed camera height)
+CAMERA_DISTANCE = 10.0  # meters (assumed camera distance from pitch)
+MAX_POSITION_JUMP = 30  # Pixels, tightened for continuous trajectory
 
 def process_video(video_path):
     if not os.path.exists(video_path):
         return [], [], [], "Error: Video file not found"
     cap = cv2.VideoCapture(video_path)
+    # Get native video resolution
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     frames = []
     ball_positions = []
-    detection_frames = []  # Track frames with detections
+    detection_frames = []
     debug_log = []
 
     frame_count = 0
@@ -34,15 +45,17 @@ def process_video(video_path):
             break
         frame_count += 1
         frames.append(frame.copy())
-        results = model.predict(frame, conf=CONF_THRESHOLD)
+        # Use native resolution for inference
+        results = model.predict(frame, conf=CONF_THRESHOLD, imgsz=(frame_height, frame_width), iou=0.5, max_det=1)
         detections = 0
         for detection in results[0].boxes:
-            if detection.cls == 0:  # Assuming class 0 is the ball
+            if detection.cls == 0:  # 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])
-                detection_frames.append(frame_count - 1)  # Store frame index (0-based)
-                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+                if detections == 1:  # Only consider frames with exactly one detection
+                    x1, y1, x2, y2 = detection.xyxy[0].cpu().numpy()
+                    ball_positions.append([(x1 + x2) / 2, (y1 + y2) / 2])
+                    detection_frames.append(frame_count - 1)
+                    cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
         frames[-1] = frame
         debug_log.append(f"Frame {frame_count}: {detections} ball detections")
     cap.release()
@@ -51,54 +64,164 @@ def process_video(video_path):
         debug_log.append("No balls detected in any frame")
     else:
         debug_log.append(f"Total ball detections: {len(ball_positions)}")
+        debug_log.append(f"Video resolution: {frame_width}x{frame_height}")
 
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
-def smooth_trajectory(ball_positions, frames):
-    if len(ball_positions) < 2:
-        return 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
-
-    # Use cubic spline interpolation to smooth the trajectory
-    try:
-        spline_x = CubicSpline(times, x_coords, bc_type='natural')
-        spline_y = CubicSpline(times, y_coords, bc_type='natural')
-    except Exception as e:
-        return None, f"Error in trajectory smoothing: {str(e)}"
-
-    # Project trajectory (detected + future for LBW decision)
-    t_full = np.linspace(times[0], times[-1] + 0.5, len(times) + 10)
-    x_full = spline_x(t_full)
-    y_full = spline_y(t_full)
-    trajectory = list(zip(x_full, y_full))
+def pixel_to_3d(x, y, frame_height, frame_width):
+    """Convert 2D pixel coordinates to 3D real-world coordinates."""
+    x_norm = x / frame_width
+    y_norm = y / frame_height
+    x_3d = (x_norm - 0.5) * 3.0  # Center x at 0 (middle of pitch)
+    y_3d = y_norm * PITCH_LENGTH
+    z_3d = (1 - y_norm) * BALL_DIAMETER * 5  # Scale to approximate ball bounce height
+    return x_3d, y_3d, z_3d
 
-    return trajectory, "Trajectory smoothed successfully"
+def estimate_trajectory(ball_positions, frames, detection_frames):
+    if len(ball_positions) < 2:
+        return None, None, None, None, None, None, None, None, None, "Error: Fewer than 2 ball detections for trajectory"
+    frame_height, frame_width = frames[0].shape[:2]
 
-def detect_pitch_and_impact(ball_positions, frames, frame_height):
+    # Filter out sudden changes in position for continuous trajectory
+    filtered_positions = [ball_positions[0]]
+    filtered_frames = [detection_frames[0]]
+    for i in range(1, len(ball_positions)):
+        prev_pos = filtered_positions[-1]
+        curr_pos = ball_positions[i]
+        distance = np.sqrt((curr_pos[0] - prev_pos[0])**2 + (curr_pos[1] - prev_pos[1])**2)
+        if distance <= MAX_POSITION_JUMP:
+            filtered_positions.append(curr_pos)
+            filtered_frames.append(detection_frames[i])
+        else:
+            # Skip sudden jumps to maintain continuity
+            continue
+
+    if len(filtered_positions) < 2:
+        return None, None, None, None, None, None, None, None, None, "Error: Fewer than 2 valid ball detections after filtering"
+
+    x_coords = [pos[0] for pos in filtered_positions]
+    y_coords = [pos[1] for pos in filtered_positions]
+    times = np.array(filtered_frames) / FRAME_RATE
+
+    # Pitch point detection: Assume it happens when the ball reaches a certain low point on the y-axis
     pitch_point = None
-    impact_point = None
-    
-    # Pitch detection: Ball hits the ground (vertical velocity and low height)
-    pitch_threshold_y = frame_height * 0.75  # Ball reaches near the ground
-
-    # For the impact point, we assume that the region of interest is near the stumps
-    impact_zone_y_min = frame_height * 0.80
-    impact_zone_y_max = frame_height * 0.85
-
-    # Detect pitch and impact points
-    for i, (x, y) in enumerate(ball_positions):
-        if y > pitch_threshold_y and not pitch_point:
-            pitch_point = (x, y)  # Ball has hit the ground
+    pitch_frame = None
+    for i in range(1, len(y_coords)):
+        if y_coords[i] > frame_height * 0.75:  # The ball reaches near the ground
+            pitch_point = filtered_positions[i]
+            pitch_frame = filtered_frames[i]
+            break
 
-        # Check if the ball is near the batsman's impact zone (e.g., near stumps)
-        if impact_zone_y_min <= y <= impact_zone_y_max and not impact_point:
-            impact_point = (x, y)  # Ball has impacted the batsman (bat or pad)
+    # Impact point detection: Look for sudden changes in the y-position (delta_y) or when ball enters impact zone
+    impact_idx = None
+    impact_frame = None
+    for i in range(1, len(y_coords)):
+        delta_y = abs(y_coords[i] - y_coords[i-1])
+        if delta_y > IMPACT_DELTA_Y:
+            impact_idx = i
+            impact_frame = filtered_frames[i]
+            break
+        elif y_coords[i] > frame_height * IMPACT_ZONE_Y:
+            impact_idx = i
+            impact_frame = filtered_frames[i]
+            break
+    if impact_idx is None:
+        impact_idx = len(filtered_positions) - 1
+        impact_frame = filtered_frames[-1]
+    impact_point = filtered_positions[impact_idx]
 
-    return pitch_point, impact_point
+    try:
+        # Use cubic interpolation for smoother trajectory
+        fx = interp1d(times[:impact_idx + 1], x_coords[:impact_idx + 1], kind='cubic', fill_value="extrapolate")
+        fy = interp1d(times[:impact_idx + 1], y_coords[:impact_idx + 1], kind='cubic', fill_value="extrapolate")
+    except Exception as e:
+        return None, None, None, None, None, None, None, None, None, f"Error in trajectory interpolation: {str(e)}"
+
+    # Generate dense points for all frames between first and last detection
+    total_frames = max(detection_frames) - min(detection_frames) + 1
+    t_full = np.linspace(times[0], times[-1], total_frames * SLOW_MOTION_FACTOR)
+    x_full = fx(t_full)
+    y_full = fy(t_full)
+    trajectory_2d = list(zip(x_full, y_full))
+
+    trajectory_3d = [pixel_to_3d(x, y, frame_height, frame_width) for x, y in trajectory_2d]
+    detections_3d = [pixel_to_3d(x, y, frame_height, frame_width) for x, y in filtered_positions]
+    pitch_point_3d = pixel_to_3d(pitch_point[0], pitch_point[1], frame_height, frame_width) if pitch_point else None
+    impact_point_3d = pixel_to_3d(impact_point[0], impact_point[1], frame_height, frame_width) if impact_point else None
+
+    debug_log = (
+        f"Trajectory estimated successfully\n"
+        f"Pitch point at frame {pitch_frame + 1}: ({pitch_point[0]:.1f}, {pitch_point[1]:.1f})\n"
+        f"Impact point at frame {impact_frame + 1}: ({impact_point[0]:.1f}, {impact_point[1]:.1f})\n"
+        f"Detections in frames: {filtered_frames}"
+    )
+    return trajectory_2d, pitch_point, impact_point, pitch_frame, impact_frame, detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, debug_log
+
+def create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, plot_type="detections"):
+    """Create 3D Plotly visualization for detections or trajectory using single-detection frames."""
+    stump_x = [-STUMPS_WIDTH/2, STUMPS_WIDTH/2, 0]
+    stump_y = [PITCH_LENGTH, PITCH_LENGTH, PITCH_LENGTH]
+    stump_z = [0, 0, 0]
+    stump_top_z = [STUMPS_HEIGHT, STUMPS_HEIGHT, STUMPS_HEIGHT]
+    bail_x = [-STUMPS_WIDTH/2, STUMPS_WIDTH/2]
+    bail_y = [PITCH_LENGTH, PITCH_LENGTH]
+    bail_z = [STUMPS_HEIGHT, STUMPS_HEIGHT]
+
+    stump_traces = []
+    for i in range(3):
+        stump_traces.append(go.Scatter3d(
+            x=[stump_x[i], stump_x[i]], y=[stump_y[i], stump_y[i]], z=[stump_z[i], stump_top_z[i]],
+            mode='lines', line=dict(color='black', width=5), name=f'Stump {i+1}'
+        ))
+    bail_traces = [
+        go.Scatter3d(
+            x=bail_x, y=bail_y, z=bail_z,
+            mode='lines', line=dict(color='black', width=5), name='Bail'
+        )
+    ]
+
+    pitch_scatter = go.Scatter3d(
+        x=[pitch_point_3d[0]] if pitch_point_3d else [], 
+        y=[pitch_point_3d[1]] if pitch_point_3d else [], 
+        z=[pitch_point_3d[2]] if pitch_point_3d else [],
+        mode='markers', marker=dict(size=8, color='red'), name='Pitch Point'
+    )
+    impact_scatter = go.Scatter3d(
+        x=[impact_point_3d[0]] if impact_point_3d else [], 
+        y=[impact_point_3d[1]] if impact_point_3d else [], 
+        z=[impact_point_3d[2]] if impact_point_3d else [],
+        mode='markers', marker=dict(size=8, color='yellow'), name='Impact Point'
+    )
+
+    if plot_type == "detections":
+        x, y, z = zip(*detections_3d) if detections_3d else ([], [], [])
+        scatter = go.Scatter3d(
+            x=x, y=y, z=z, mode='markers',
+            marker=dict(size=5, color='green'), name='Single Ball Detections'
+        )
+        data = [scatter, pitch_scatter, impact_scatter] + stump_traces + bail_traces
+        title = "3D Single Ball Detections"
+    else:
+        x, y, z = zip(*trajectory_3d) if trajectory_3d else ([], [], [])
+        trajectory_line = go.Scatter3d(
+            x=x, y=y, z=z, mode='lines',
+            line=dict(color='blue', width=4), name='Ball Trajectory (Single Detections)'
+        )
+        data = [trajectory_line, pitch_scatter, impact_scatter] + stump_traces + bail_traces
+        title = "3D Ball Trajectory (Single Detections)"
+
+    layout = go.Layout(
+        title=title,
+        scene=dict(
+            xaxis_title='X (meters)', yaxis_title='Y (meters)', zaxis_title='Z (meters)',
+            xaxis=dict(range=[-1.5, 1.5]), yaxis=dict(range=[0, PITCH_LENGTH]),
+            zaxis=dict(range=[0, STUMPS_HEIGHT * 2]), aspectmode='manual',
+            aspectratio=dict(x=1, y=4, z=0.5)
+        ),
+        showlegend=True
+    )
+    fig = go.Figure(data=data, layout=layout)
+    return fig
 
 def lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point):
     if not frames:
@@ -108,105 +231,101 @@ def lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point):
 
     frame_height, frame_width = frames[0].shape[:2]
     stumps_x = frame_width / 2
-    stumps_y = frame_height * 0.9  # Position of the stumps at the bottom of the frame
+    stumps_y = frame_height * 0.9
     stumps_width_pixels = frame_width * (STUMPS_WIDTH / 3.0)
 
     pitch_x, pitch_y = pitch_point
     impact_x, impact_y = impact_point
 
-    # Check pitching point - the ball should land between stumps
     if pitch_x < stumps_x - stumps_width_pixels / 2 or pitch_x > stumps_x + stumps_width_pixels / 2:
         return f"Not Out (Pitched outside line at x: {pitch_x:.1f}, y: {pitch_y:.1f})", trajectory, pitch_point, impact_point
-
-    # Check impact point - the ball should hit within the stumps area
     if impact_x < stumps_x - stumps_width_pixels / 2 or impact_x > stumps_x + stumps_width_pixels / 2:
         return f"Not Out (Impact outside line at x: {impact_x:.1f}, y: {impact_y:.1f})", trajectory, pitch_point, impact_point
-
-    # 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 f"Out (Ball hits stumps, Pitch at x: {pitch_x:.1f}, y: {pitch_y:.1f}, Impact at x: {impact_x:.1f}, y: {impact_y:.1f})", trajectory, pitch_point, impact_point
-
     return f"Not Out (Missing stumps, Pitch at x: {pitch_x:.1f}, y: {pitch_y:.1f}, Impact at x: {impact_x:.1f}, y: {impact_y:.1f})", trajectory, pitch_point, impact_point
 
-def generate_slow_motion(frames, trajectory, pitch_point, impact_point, detection_frames, output_path):
+def generate_slow_motion(frames, trajectory, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path):
     if not frames:
         return None
+    frame_height, frame_width = frames[0].shape[:2]
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frames[0].shape[1], frames[0].shape[0]))
-
-    trajectory_points = np.array(trajectory[:len(detection_frames)], dtype=np.int32).reshape((-1, 1, 2))
-
-    pitch_point_detected = False
-    impact_point_detected = False
+    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frame_width, frame_height))
+
+    # Map trajectory points to all frames between first and last detection
+    if trajectory and detection_frames:
+        min_frame = min(detection_frames)
+        max_frame = max(detection_frames)
+        total_frames = max_frame - min_frame + 1
+        trajectory_points = np.array(trajectory, dtype=np.int32).reshape((-1, 1, 2))
+        traj_per_frame = len(trajectory) // total_frames
+        trajectory_indices = [i * traj_per_frame for i in range(total_frames)]
+    else:
+        trajectory_points = np.array([], dtype=np.int32)
+        trajectory_indices = []
 
     for i, frame in enumerate(frames):
-        # Draw trajectory (blue line) only for frames with detections
-        if i in detection_frames and trajectory_points.size > 0:
-            cv2.polylines(frame, [trajectory_points[:detection_frames.index(i) + 1]], False, (255, 0, 0), 2)
-
-        # Draw pitch point (red circle with label) when the ball touches the ground
-        if pitch_point and not pitch_point_detected:
+        frame_idx = i - min_frame if trajectory_indices else -1
+        if frame_idx >= 0 and frame_idx < total_frames and trajectory_points.size > 0:
+            # Draw trajectory up to current frame
+            end_idx = trajectory_indices[frame_idx] + 1
+            cv2.polylines(frame, [trajectory_points[:end_idx]], False, (255, 0, 0), 2)
+        if pitch_point and i == pitch_frame:
             x, y = pitch_point
-            if y > frame.shape[0] * 0.75:  # Adjust this threshold for the ground position
-                pitch_point_detected = True
-        if pitch_point_detected:
             cv2.circle(frame, (int(x), int(y)), 8, (0, 0, 255), -1)
             cv2.putText(frame, "Pitch Point", (int(x) + 10, int(y) - 10), 
                         cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
-
-        # Draw impact point (yellow circle with label) when ball is near stumps
-        if impact_point and not impact_point_detected:
+        if impact_point and i == impact_frame:
             x, y = impact_point
-            if y > frame.shape[0] * 0.85:  # Adjust this threshold for impact point
-                impact_point_detected = True
-        if impact_point_detected:
             cv2.circle(frame, (int(x), int(y)), 8, (0, 255, 255), -1)
             cv2.putText(frame, "Impact Point", (int(x) + 10, int(y) + 20), 
                         cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
-
-        # Add wicket lines for the stumps
-        stumps_x = frame.shape[1] // 2
-        stumps_y = frame.shape[0] * 0.9
-        stumps_width = frame.shape[1] * 0.1
-        cv2.line(frame, (int(stumps_x - stumps_width / 2), int(stumps_y)),
-                 (int(stumps_x + stumps_width / 2), int(stumps_y)), (0, 255, 0), 3)
-
-        # Write frames to output video
         for _ in range(SLOW_MOTION_FACTOR):
             out.write(frame)
-
     out.release()
     return output_path
 
 def drs_review(video):
     frames, ball_positions, detection_frames, debug_log = process_video(video)
     if not frames:
-        return f"Error: Failed to process video", None
-    
-    trajectory, smoothing_log = smooth_trajectory(ball_positions, frames)
+        return f"Error: Failed to process video\nDebug Log:\n{debug_log}", None, None, None
 
-    # Detect pitch and impact points based on ball positions
-    pitch_point, impact_point = detect_pitch_and_impact(ball_positions, frames, frames[0].shape[0])
+    trajectory_2d, pitch_point, impact_point, pitch_frame, impact_frame, detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, trajectory_log = estimate_trajectory(ball_positions, frames, detection_frames)
+    
+    if trajectory_2d is None:
+        return (f"Error: {trajectory_log}\nDebug Log:\n{debug_log}", None, None, None)
 
-    decision, trajectory, pitch_point, impact_point = lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point)
+    decision, trajectory_2d, pitch_point, impact_point = lbw_decision(ball_positions, trajectory_2d, frames, pitch_point, impact_point)
 
     output_path = f"output_{uuid.uuid4()}.mp4"
-    slow_motion_path = generate_slow_motion(frames, trajectory, pitch_point, impact_point, detection_frames, output_path)
+    slow_motion_path = generate_slow_motion(frames, trajectory_2d, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path)
+
+    detections_fig = None
+    trajectory_fig = None
+    if detections_3d:
+        detections_fig = create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, "detections")
+        trajectory_fig = create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, "trajectory")
 
-    return f"DRS Decision: {decision}", slow_motion_path
+    debug_output = f"{debug_log}\n{trajectory_log}"
+    return (f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", 
+            slow_motion_path, 
+            detections_fig, 
+            trajectory_fig)
 
 # Gradio interface
 iface = gr.Interface(
     fn=drs_review,
     inputs=gr.Video(label="Upload Video Clip"),
     outputs=[
-        gr.Textbox(label="DRS Decision"),
-        gr.Video(label="Slow-Motion Replay with Ball Detection (Green), Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow), Wicket Lines")
+        gr.Textbox(label="DRS Decision and Debug Log"),
+        gr.Video(label="Very Slow-Motion Replay with Ball Detection (Green), Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow)"),
+        gr.Plot(label="3D Single Ball Detections Plot"),
+        gr.Plot(label="3D Ball Trajectory Plot (Single Detections)")
     ],
     title="AI-Powered DRS for LBW in Local Cricket",
-    description="Upload a video clip of a cricket delivery to get an LBW decision and slow-motion replay showing ball detection (green boxes), trajectory (blue line), pitch point (red circle), impact point (yellow circle), and wicket lines."
+    description="Upload a video clip of a cricket delivery to get an LBW decision, a slow-motion replay, and 3D visualizations. The replay shows ball detection (green boxes), trajectory (blue line), pitch point (red circle), and impact point (yellow circle). The 3D plots show single-detection frames (green markers) and trajectory (blue line) with wicket lines (black), pitch point (red), and impact point (yellow)."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()