Update app.py

app.py

@@ -4,6 +4,7 @@ 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
 
@@ -13,13 +14,15 @@ 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 = 20  # Input video frame rate
-SLOW_MOTION_FACTOR = 3  # Adjusted for 20 FPS
+FRAME_RATE = 30  # Input video frame rate
+SLOW_MOTION_FACTOR = 6  # For very slow motion (6x slower)
 CONF_THRESHOLD = 0.25  # Confidence threshold for detection
-IMPACT_ZONE_Y = 0.80  # Fraction of frame height for impact zone
-PITCH_ZONE_Y = 0.80  # Fraction of frame height for pitch zone
-IMPACT_DELTA_Y = 75  # Pixels for detecting sudden y-position change
-IMPACT_DELTA_X = 50  # Pixels for x-position stability
+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)
 
 def process_video(video_path):
     if not os.path.exists(video_path):
@@ -27,7 +30,7 @@ def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     frames = []
     ball_positions = []
-    detection_frames = []  # Track frames with exactly one detection
+    detection_frames = []
     debug_log = []
 
     frame_count = 0
@@ -38,96 +41,82 @@ def process_video(video_path):
         frame_count += 1
         frames.append(frame.copy())
         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()
-            x_center = (x1 + x2) / 2
-            y_center = (y1 + y2) / 2
-            ball_positions.append([x_center, y_center])
-            detection_frames.append(frame_count - 1)  # 0-based index
-            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
-            debug_log.append(f"Frame {frame_count}: 1 ball detection at (x: {x_center:.1f}, y: {y_center:.1f})")
-        else:
-            debug_log.append(f"Frame {frame_count}: {len(detections)} ball detections (ignored)")
+        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)
         frames[-1] = frame
+        debug_log.append(f"Frame {frame_count}: {detections} ball detections")
     cap.release()
 
     if not ball_positions:
-        debug_log.append("No valid single-ball detections in any frame")
+        debug_log.append("No balls detected in any frame")
     else:
-        debug_log.append(f"Total valid single-ball detections: {len(ball_positions)}")
+        debug_log.append(f"Total ball detections: {len(ball_positions)}")
 
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
-def estimate_trajectory(ball_positions, detection_frames, frames):
+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):
     if len(ball_positions) < 2:
-        return None, None, None, None, None, None, "Error: Fewer than 2 valid single-ball detections for trajectory"
-    frame_height = frames[0].shape[0]
+        return None, None, None, None, None, "Error: Fewer than 2 ball detections for trajectory"
+    frame_height, frame_width = frames[0].shape[:2]
 
-    # 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: first detection where 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
-    else:
-        debug_log = "Warning: No pitch point detected (y never exceeds PITCH_ZONE_Y), using first detection"
-        pitch_idx = 0
-    pitch_point = ball_positions[pitch_idx]
-    pitch_frame = detection_frames[pitch_idx]
+    pitch_point = ball_positions[0]
+    pitch_frame = detection_frames[0]
 
-    # Impact point: sudden y-change, x stability, or y exceeds IMPACT_ZONE_Y
     impact_idx = None
+    impact_frame = None
     for i in range(1, len(y_coords)):
-        delta_y = abs(y_coords[i] - y_coords[i-1])
-        delta_x = abs(x_coords[i] - x_coords[i-1])
-        if (y_coords[i] > frame_height * IMPACT_ZONE_Y or
-            (delta_y > IMPACT_DELTA_Y and delta_x < IMPACT_DELTA_X)):
+        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
-        debug_log = f"Warning: No impact point detected (no sudden y-change or y > IMPACT_ZONE_Y), using last detection (frame {detection_frames[impact_idx] + 1})"
-    else:
-        debug_log = ""
+        impact_frame = detection_frames[-1]
     impact_point = ball_positions[impact_idx]
-    impact_frame = detection_frames[impact_idx]
-
-    # Use only detected positions up to impact 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, x_coords, kind='linear', fill_value="extrapolate")
-        fy = interp1d(times, y_coords, kind='quadratic', fill_value="extrapolate")
+        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")
     except Exception as 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))
+        return None, None, None, None, None, f"Error in trajectory interpolation: {str(e)}"
 
-    # 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)
-    full_trajectory = list(zip(x_full, y_full))
+    trajectory_2d = list(zip(x_full, y_full))
 
-    debug_log = (f"{debug_log}\nTrajectory 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
+    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)
 
-def lbw_decision(ball_positions, full_trajectory, frames, pitch_point, impact_point):
+    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):
     if not frames:
         return "Error: No frames processed", None, None, None
-    if not full_trajectory or len(ball_positions) < 2:
-        return "Not enough data (insufficient valid single-ball detections)", None, None, None
+    if not trajectory or len(ball_positions) < 2:
+        return "Not enough data (insufficient ball detections)", None, None, None
 
     frame_height, frame_width = frames[0].shape[:2]
     stumps_x = frame_width / 2
@@ -137,52 +126,111 @@ def lbw_decision(ball_positions, full_trajectory, frames, pitch_point, impact_po
     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})", full_trajectory, pitch_point, impact_point
-
-    # Check impact point
+        return f"Not Out (Pitched outside line at x: {pitch_x:.1f}, y: {pitch_y:.1f})", trajectory, pitch_point, 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})", full_trajectory, pitch_point, impact_point
-
-    # Check trajectory hitting stumps
-    for x, y in full_trajectory:
+        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:
         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})", 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):
+            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."""
+    # Wicket lines (stumps and bails)
+    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]
+
+    # Stumps (three vertical lines)
+    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}'
+        ))
+    # Bails (two horizontal lines)
+    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":
+        # Ball detections plot
+        x, y, z = zip(*detections_3d)
+        scatter = go.Scatter3d(
+            x=x, y=y, z=z, mode='markers',
+            marker=dict(size=5, color='green'), name='Ball Detections'
+        )
+        # Pitch and impact points
+        pitch_scatter = go.Scatter3d(
+            x=[pitch_point_3d[0]], y=[pitch_point_3d[1]], z=[pitch_point_3d[2]],
+            mode='markers', marker=dict(size=8, color='red'), name='Pitch Point'
+        )
+        impact_scatter = go.Scatter3d(
+            x=[impact_point_3d[0]], y=[impact_point_3d[1]], z=[impact_point_3d[2]],
+            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:
+        # Trajectory plot
+        x, y, z = zip(*trajectory_3d)
+        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]], y=[pitch_point_3d[1]], z=[pitch_point_3d[2]],
+            mode='markers', marker=dict(size=8, color='red'), name='Pitch Point'
+        )
+        impact_scatter = go.Scatter3d(
+            x=[impact_point_3d[0]], y=[impact_point_3d[1]], z=[impact_point_3d[2]],
+            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):
     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, (frame_width, frame_height))
+    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frames[0].shape[1], frames[0].shape[0]))
 
-    # Prepare trajectory points for visualization
-    trajectory_points = np.array(vis_trajectory, dtype=np.int32).reshape((-1, 1, 2))
+    trajectory_points = np.array(trajectory[:len(detection_frames)], 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:
-            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
+            cv2.polylines(frame, [trajectory_points[:detection_frames.index(i) + 1]], False, (255, 0, 0), 2)
         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()
@@ -191,15 +239,29 @@ def generate_slow_motion(frames, vis_trajectory, pitch_point, pitch_frame, impac
 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
-    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)
+        return f"Error: Failed to process video\nDebug Log:\n{debug_log}", None, None, None
+    trajectory, 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)
+    decision, trajectory, pitch_point, impact_point = lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point)
 
     output_path = f"output_{uuid.uuid4()}.mp4"
-    slow_motion_path = generate_slow_motion(frames, vis_trajectory, pitch_point, pitch_frame, impact_point, impact_frame, detection_frames, output_path)
+    slow_motion_path = generate_slow_motion(frames, trajectory, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path)
+
+    detections_plot_path = f"detections_3d_{uuid.uuid4()}.html"
+    trajectory_plot_path = f"trajectory_3d_{uuid.uuid4()}.html"
+    if detections_3d:
+        detections_fig = create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, "detections")
+        detections_fig.write_html(detections_plot_path)
+        trajectory_fig = create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d, "trajectory")
+        trajectory_fig.write_html(trajectory_plot_path)
+    else:
+        detections_plot_path = None
+        trajectory_plot_path = None
 
     debug_output = f"{debug_log}\n{trajectory_log}"
-    return f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", slow_motion_path
+    return (f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", 
+            slow_motion_path, 
+            detections_plot_path, 
+            trajectory_plot_path)
 
 # Gradio interface
 iface = gr.Interface(
@@ -207,10 +269,12 @@ 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.Video(label="Very Slow-Motion Replay with Ball Detection (Green), Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow)"),
+        gr.File(label="3D Ball Detections Plot (HTML)"),
+        gr.File(label="3D Ball Trajectory Plot (HTML)")
     ],
     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), and impact point (yellow circle)."
+    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 and trajectory with wicket lines."
 )
 
 if __name__ == "__main__":