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
 
@@ -14,24 +15,25 @@ 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
+FRAME_RATE = 20  # Input video frame rate (to be updated dynamically)
 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
+CONF_THRESHOLD = 0.4  # Increased confidence threshold for better detection
+IMPACT_ZONE_Y = 0.8  # Adjusted fraction of frame height for impact zone
+IMPACT_VELOCITY_THRESHOLD = 1000  # Pixels/second for detecting impact
 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
+MAX_POSITION_JUMP = 50  # Increased for smoother 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
+    # Get native video resolution and frame rate
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    FRAME_RATE = cap.get(cv2.CAP_PROP_FPS) or 20  # Use actual frame rate or default
     frames = []
     ball_positions = []
     detection_frames = []
@@ -44,8 +46,9 @@ def process_video(video_path):
             break
         frame_count += 1
         frames.append(frame.copy())
-        # Use smaller image size to speed up detection
-        results = model.predict(frame, conf=CONF_THRESHOLD, imgsz=(384, 640), iou=0.5, max_det=1)
+        # Enhance frame contrast for better detection
+        frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=10)
+        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:  # Class 0 is the ball
@@ -57,6 +60,8 @@ def process_video(video_path):
                     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")
+        # Save debug frame
+        cv2.imwrite(f"debug_frame_{frame_count}.jpg", frame)
     cap.release()
 
     if not ball_positions:
@@ -64,6 +69,7 @@ def process_video(video_path):
     else:
         debug_log.append(f"Total ball detections: {len(ball_positions)}")
         debug_log.append(f"Video resolution: {frame_width}x{frame_height}")
+        debug_log.append(f"Video frame rate: {FRAME_RATE}")
 
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
@@ -92,7 +98,6 @@ def estimate_trajectory(ball_positions, frames, detection_frames):
             filtered_positions.append(curr_pos)
             filtered_frames.append(detection_frames[i])
         else:
-            # Skip sudden jumps to maintain continuity
             continue
 
     if len(filtered_positions) < 2:
@@ -102,21 +107,21 @@ def estimate_trajectory(ball_positions, frames, detection_frames):
     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
-    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
+    # Convert to 3D for pitch point detection
+    detections_3d = [pixel_to_3d(x, y, frame_height, frame_width) for x, y in filtered_positions]
+    
+    # Pitch point: Detection with lowest z-coordinate (closest to ground)
+    pitch_idx = min(range(len(detections_3d)), key=lambda i: detections_3d[i][2])
+    pitch_point = filtered_positions[pitch_idx]
+    pitch_frame = filtered_frames[pitch_idx]
 
-    # Impact point detection: Look for sudden changes in the y-position (delta_y) or when ball enters impact zone
+    # Impact point: Detect sudden velocity change or impact zone
     impact_idx = None
     impact_frame = None
+    velocities = [np.sqrt((x_coords[i] - x_coords[i-1])**2 + (y_coords[i] - y_coords[i-1])**2) / (times[i] - times[i-1]) 
+                  for i in range(1, len(x_coords))]
     for i in range(1, len(y_coords)):
-        delta_y = abs(y_coords[i] - y_coords[i-1])
-        if delta_y > IMPACT_DELTA_Y:
+        if velocities[i-1] > IMPACT_VELOCITY_THRESHOLD:
             impact_idx = i
             impact_frame = filtered_frames[i]
             break
@@ -130,52 +135,113 @@ def estimate_trajectory(ball_positions, frames, detection_frames):
     impact_point = filtered_positions[impact_idx]
 
     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")
+        # Use linear interpolation for more stable trajectory
+        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='linear', 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)
+    t_full = np.linspace(times[0], times[impact_idx], 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]
-
-    # Handle missing pitch and impact points gracefully
-    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
-
-    # Handle cases where no pitch/impact point is found
-    if pitch_point is None:
-        pitch_frame = "N/A"
-        pitch_point_3d = None  # No 3D coordinates for pitch point
-    if impact_point is None:
-        impact_frame = "N/A"
-        impact_point_3d = None  # No 3D coordinates for impact point
+    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 trajectory and points
     debug_log = (
         f"Trajectory estimated successfully\n"
-        f"Pitch point at frame {pitch_frame + 1 if pitch_frame != 'N/A' else 'N/A'}: {pitch_point if pitch_point else 'Not detected'}\n"
-        f"Impact point at frame {impact_frame + 1 if impact_frame != 'N/A' else 'N/A'}: {impact_point if impact_point else 'Not detected'}\n"
-        f"Detections in frames: {filtered_frames}"
+        f"Pitch point at frame {pitch_frame + 1}: ({pitch_point[0]:.1f}, {pitch_point[1]:.1f}), 3D: {pitch_point_3d}\n"
+        f"Impact point at frame {impact_frame + 1}: ({impact_point[0]:.1f}, {impact_point[1]:.1f}), 3D: {impact_point_3d}\n"
+        f"Detections in frames: {filtered_frames}\n"
+        f"Velocities: {velocities}"
     )
+    # Save trajectory plot for debugging
+    import matplotlib.pyplot as plt
+    plt.plot(x_coords, y_coords, 'bo-', label='Filtered Detections')
+    plt.plot(pitch_point[0], pitch_point[1], 'ro', label='Pitch Point')
+    plt.plot(impact_point[0], impact_point[1], 'yo', label='Impact Point')
+    plt.legend()
+    plt.savefig("trajectory_debug.png")
+
     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]] ifpitch_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:
         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
 
-    # Check for None values before unpacking
-    if pitch_point is None or impact_point is None:
-        return "Not Out (Unable to determine pitch or impact points)", trajectory, pitch_point, impact_point
-
     frame_height, frame_width = frames[0].shape[:2]
     stumps_x = frame_width / 2
     stumps_y = frame_height * 0.9
@@ -200,7 +266,6 @@ def generate_slow_motion(frames, trajectory, pitch_point, impact_point, detectio
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     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)
@@ -215,7 +280,6 @@ def generate_slow_motion(frames, trajectory, pitch_point, impact_point, detectio
     for i, frame in enumerate(frames):
         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:
@@ -248,9 +312,17 @@ def drs_review(video):
     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")
+
     debug_output = f"{debug_log}\n{trajectory_log}"
     return (f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", 
-            slow_motion_path)
+            slow_motion_path, 
+            detections_fig, 
+            trajectory_fig)
 
 # Gradio interface
 iface = gr.Interface(
@@ -259,10 +331,12 @@ iface = gr.Interface(
     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 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, 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)."
+    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()