Update app.py

app.py

@@ -4,7 +4,6 @@ import torch
 from ultralytics import YOLO
 import gradio as gr
 from scipy.interpolate import interp1d
-import plotly.graph_objects as go
 import uuid
 import os
 
@@ -14,15 +13,12 @@ model = YOLO("best.pt")
 # 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.2  # Lowered confidence threshold to improve detection
-IMPACT_ZONE_Y = 0.85  # Fraction of frame height where impact is likely
+FRAME_RATE = 20  # Input video frame rate
+SLOW_MOTION_FACTOR = 3  # Adjusted for 20 FPS
+CONF_THRESHOLD = 0.25  # Confidence threshold for detection
+IMPACT_ZONE_Y = 0.85  # Fraction of frame height for impact zone
+PITCH_ZONE_Y = 0.75  # Fraction of frame height for pitch zone
 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)
 
 def process_video(video_path):
     if not os.path.exists(video_path):
@@ -30,7 +26,7 @@ def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     frames = []
     ball_positions = []
-    detection_frames = []
+    detection_frames = []  # Track frames with exactly one detection
     debug_log = []
 
     frame_count = 0
@@ -40,83 +36,85 @@ def process_video(video_path):
             break
         frame_count += 1
         frames.append(frame.copy())
-        results = model.predict(frame, conf=CONF_THRESHOLD, imgsz=640)
-        detections = 0
-        for detection in results[0].boxes:
-            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)
-                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+        results = model.predict(frame, conf=CONF_THRESHOLD)
+        detections = [det for det in results[0].boxes if det.cls == 0]  # Class 0 is cricketBall
+        if len(detections) == 1:  # Only consider frames with exactly one detection
+            x1, y1, x2, y2 = detections[0].xyxy[0].cpu().numpy()
+            ball_positions.append([(x1 + x2) / 2, (y1 + y2) / 2])
+            detection_frames.append(frame_count - 1)  # 0-based index
+            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")
+        debug_log.append(f"Frame {frame_count}: {len(detections)} ball detections")
     cap.release()
 
     if not ball_positions:
-        debug_log.append("No balls detected in any frame")
+        debug_log.append("No valid single-ball detections in any frame")
     else:
-        debug_log.append(f"Total ball detections: {len(ball_positions)}")
+        debug_log.append(f"Total valid single-ball detections: {len(ball_positions)}")
 
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
-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
-
-def estimate_trajectory(ball_positions, frames, detection_frames):
+def estimate_trajectory(ball_positions, detection_frames, 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]
+        return None, None, None, None, None, None, "Error: Fewer than 2 valid single-ball detections for trajectory"
+    frame_height = frames[0].shape[0]
 
+    # Extract x, y coordinates
     x_coords = [pos[0] for pos in ball_positions]
     y_coords = [pos[1] for pos in ball_positions]
     times = np.array(detection_frames) / FRAME_RATE
 
-    pitch_point = ball_positions[0]
-    pitch_frame = detection_frames[0]
+    # Pitch point: first valid detection or when y exceeds PITCH_ZONE_Y
+    pitch_idx = 0
+    for i, y in enumerate(y_coords):
+        if y > frame_height * PITCH_ZONE_Y:
+            pitch_idx = i
+            break
+    pitch_point = ball_positions[pitch_idx]
+    pitch_frame = detection_frames[pitch_idx]
 
+    # Impact point: sudden y-change or y exceeds IMPACT_ZONE_Y
     impact_idx = None
-    impact_frame = None
     for i in range(1, len(y_coords)):
-        if y_coords[i] > frame_height * IMPACT_ZONE_Y or abs(y_coords[i] - y_coords[i-1]) > IMPACT_DELTA_Y:
+        if (y_coords[i] > frame_height * IMPACT_ZONE_Y or
+            abs(y_coords[i] - y_coords[i-1]) > IMPACT_DELTA_Y):
             impact_idx = i
-            impact_frame = detection_frames[i]
             break
     if impact_idx is None:
         impact_idx = len(ball_positions) - 1
-        impact_frame = detection_frames[-1]
     impact_point = ball_positions[impact_idx]
+    impact_frame = detection_frames[impact_idx]
+
+    # Use only detected positions for trajectory
+    x_coords = x_coords[:impact_idx + 1]
+    y_coords = y_coords[:impact_idx + 1]
+    times = times[:impact_idx + 1]
 
     try:
-        fx = interp1d(times[:impact_idx + 1], x_coords[:impact_idx + 1], kind='linear', fill_value="extrapolate")
-        fy = interp1d(times[:impact_idx + 1], y_coords[:impact_idx + 1], kind='quadratic', fill_value="extrapolate")
+        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, None, None, None, None, None, None, None, f"Error in trajectory interpolation: {str(e)}"
+        return None, None, None, None, None, None, f"Error in trajectory interpolation: {str(e)}"
+
+    # Trajectory for visualization (detected frames only)
+    vis_trajectory = list(zip(x_coords, y_coords))
 
+    # Full trajectory for LBW (includes projection)
     t_full = np.linspace(times[0], times[-1] + 0.5, len(times) + 10)
     x_full = fx(t_full)
     y_full = fy(t_full)
-    trajectory_2d = list(zip(x_full, y_full))
+    full_trajectory = 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 ball_positions]
-    pitch_point_3d = pixel_to_3d(pitch_point[0], pitch_point[1], frame_height, frame_width)
-    impact_point_3d = pixel_to_3d(impact_point[0], impact_point[1], frame_height, frame_width)
+    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})")
+    return full_trajectory, vis_trajectory, pitch_point, pitch_frame, impact_point, impact_frame, debug_log
 
-    debug_log = f"Trajectory estimated successfully\nPitch point at frame {pitch_frame + 1}: ({pitch_point[0]:.1f}, {pitch_point[1]:.1f})\nImpact point at frame {impact_frame + 1}: ({impact_point[0]:.1f}, {impact_point[1]:.1f})"
-    return trajectory_2d, pitch_point, impact_point, pitch_frame, impact_frame, detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, debug_log
-
-def lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point):
+def lbw_decision(ball_positions, full_trajectory, frames, pitch_point, impact_point):
     if not frames:
         return "Error: No frames processed", None, None, None
-    if not trajectory or len(ball_positions) < 2:
-        return "Not enough data (insufficient ball detections)", None, None, None
+    if not full_trajectory or len(ball_positions) < 2:
+        return "Not enough data (insufficient valid single-ball detections)", None, None, None
 
     frame_height, frame_width = frames[0].shape[:2]
     stumps_x = frame_width / 2
@@ -126,116 +124,52 @@ def lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point):
     pitch_x, pitch_y = pitch_point
     impact_x, impact_y = impact_point
 
+    # Check pitching point
     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
+        return f"Not Out (Pitched outside line at x: {pitch_x:.1f}, y: {pitch_y:.1f})", full_trajectory, pitch_point, impact_point
+
+    # Check impact point
     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
-    for x, y in trajectory:
+        return f"Not Out (Impact outside line at x: {impact_x:.1f}, y: {impact_y:.1f})", full_trajectory, pitch_point, impact_point
+
+    # Check trajectory hitting stumps
+    for x, y in full_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 create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, plot_type="detections"):
-    """Create 3D Plotly visualization for detections or trajectory."""
-    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'
-        )
-    ]
-
-    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='Ball Detections'
-        )
-        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'
-        )
-        data = [scatter, pitch_scatter, impact_scatter] + stump_traces + bail_traces
-        title = "3D 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'
-        )
-        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'
-        )
-        data = [trajectory_line, pitch_scatter, impact_scatter] + stump_traces + bail_traces
-        title = "3D Ball Trajectory"
-
-    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 generate_slow_motion(frames, trajectory, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path):
+            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})", full_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})", full_trajectory, pitch_point, impact_point
+
+def generate_slow_motion(frames, vis_trajectory, pitch_point, pitch_frame, impact_point, impact_frame, detection_frames, 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]))
+    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frame_width, frame_height))
 
-    if trajectory and detection_frames:
-        trajectory_points = np.array(trajectory[:len(detection_frames)], dtype=np.int32).reshape((-1, 1, 2))
-    else:
-        trajectory_points = np.array([], dtype=np.int32)
+    # Prepare trajectory points for visualization
+    trajectory_points = np.array(vis_trajectory, dtype=np.int32).reshape((-1, 1, 2))
 
     for i, frame in enumerate(frames):
+        # Draw trajectory (blue line) only for detected frames
         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)
+            idx = detection_frames.index(i) + 1
+            if idx <= len(trajectory_points):
+                cv2.polylines(frame, [trajectory_points[:idx]], False, (255, 0, 0), 2)
+
+        # Draw pitch point (red circle) only in pitch frame
         if pitch_point and i == pitch_frame:
             x, y = pitch_point
             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) only in impact frame
         if impact_point and i == impact_frame:
             x, y = impact_point
             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)
+
         for _ in range(SLOW_MOTION_FACTOR):
             out.write(frame)
     out.release()
@@ -244,29 +178,15 @@ def generate_slow_motion(frames, trajectory, pitch_point, impact_point, detectio
 def drs_review(video):
     frames, ball_positions, detection_frames, debug_log = process_video(video)
     if not frames:
-        return f"Error: Failed to process video\nDebug Log:\n{debug_log}", None, None, None
-
-    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_2d, pitch_point, impact_point = lbw_decision(ball_positions, trajectory_2d, frames, pitch_point, impact_point)
+        return f"Error: Failed to process video\nDebug Log:\n{debug_log}", None
+    full_trajectory, vis_trajectory, pitch_point, pitch_frame, impact_point, impact_frame, trajectory_log = estimate_trajectory(ball_positions, detection_frames, frames)
+    decision, full_trajectory, pitch_point, impact_point = lbw_decision(ball_positions, full_trajectory, frames, pitch_point, impact_point)
 
     output_path = f"output_{uuid.uuid4()}.mp4"
-    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")
+    slow_motion_path = generate_slow_motion(frames, vis_trajectory, pitch_point, pitch_frame, impact_point, impact_frame, detection_frames, output_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)
+    return f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", slow_motion_path
 
 # Gradio interface
 iface = gr.Interface(
@@ -274,12 +194,10 @@ iface = gr.Interface(
     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 (Green), Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow)"),
-        gr.Plot(label="3D Ball Detections Plot"),
-        gr.Plot(label="3D Ball Trajectory Plot")
+        gr.Video(label="Very Slow-Motion Replay with Ball Detection (Green), Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow)")
     ],
     title="AI-Powered DRS for LBW in Local Cricket",
-    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 detections (green markers) and trajectory (blue line) with wicket lines (black), pitch point (red), and impact point (yellow)."
+    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), and impact point (yellow circle)."
 )
 
 if __name__ == "__main__":