Update app.py

app.py

@@ -7,6 +7,7 @@ from scipy.interpolate import interp1d
 import plotly.graph_objects as go
 import uuid
 import os
+import plotly.io as pio
 
 # Load the trained YOLOv8n model with optimizations
 model = YOLO("best.pt")
@@ -92,9 +93,6 @@ def estimate_trajectory(ball_positions, frames, detection_frames):
         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"
@@ -106,7 +104,7 @@ def estimate_trajectory(ball_positions, frames, detection_frames):
     pitch_point = filtered_positions[0]
     pitch_frame = filtered_frames[0]
 
-    # Prioritize sudden y-change for impact detection
+    # Detect impact point
     impact_idx = None
     impact_frame = None
     for i in range(1, len(y_coords)):
@@ -125,17 +123,16 @@ 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")
+        # Cubic interpolation for smooth trajectory
+        fx = interp1d(times, x_coords, kind='cubic', fill_value="extrapolate")
+        fy = interp1d(times, y_coords, kind='cubic', fill_value="extrapolate")
+        t_full = np.linspace(times[0], times[-1], len(times) * 4)
+        x_full = fx(t_full)
+        y_full = fy(t_full)
+        trajectory_2d = list(zip(x_full, y_full))
     except Exception as e:
         return None, None, None, None, None, None, None, None, None, f"Error in trajectory interpolation: {str(e)}"
 
-    t_full = np.linspace(times[0], times[-1], len(times) * 4)  # Dense points for smooth trajectory
-    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)
@@ -249,6 +246,89 @@ def create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d
     fig = go.Figure(data=data, layout=layout)
     return fig
 
+def generate_3d_trajectory_video(trajectory_3d, pitch_point_3d, impact_point_3d, detection_frames, pitch_frame, impact_frame, output_path):
+    if not trajectory_3d or not detection_frames:
+        return None
+
+    # Define video parameters
+    frame_width = 1280
+    frame_height = 720
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frame_width, frame_height))
+
+    # 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]
+
+    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'
+        )
+    ]
+
+    # Generate frames for each detection
+    for i, frame_idx in enumerate(detection_frames):
+        # Trajectory up to current frame
+        traj_idx = min(i * 4, len(trajectory_3d) - 1)  # Match 2D trajectory density
+        x, y, z = zip(*trajectory_3d[:traj_idx + 1]) 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 point (red marker) if at or after pitch_frame
+        pitch_scatter = go.Scatter3d(
+            x=[pitch_point_3d[0]] if frame_idx >= pitch_frame else [],
+            y=[pitch_point_3d[1]] if frame_idx >= pitch_frame else [],
+            z=[pitch_point_3d[2]] if frame_idx >= pitch_frame else [],
+            mode='markers', marker=dict(size=8, color='red'), name='Pitch Point'
+        )
+
+        # Impact point (yellow marker) if at or after impact_frame
+        impact_scatter = go.Scatter3d(
+            x=[impact_point_3d[0]] if frame_idx >= impact_frame else [],
+            y=[impact_point_3d[1]] if frame_idx >= impact_frame else [],
+            z=[impact_point_3d[2]] if frame_idx >= impact_frame else [],
+            mode='markers', marker=dict(size=8, color='yellow'), name='Impact Point'
+        )
+
+        # Create frame
+        fig = go.Figure(
+            data=[trajectory_line, pitch_scatter, impact_scatter] + stump_traces + bail_traces,
+            layout=go.Layout(
+                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
+            )
+        )
+
+        # Render frame to image
+        img_bytes = pio.to_image(fig, format='png', width=frame_width, height=frame_height)
+        img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
+
+        # Write frame multiple times for slow motion
+        for _ in range(SLOW_MOTION_FACTOR):
+            out.write(img)
+
+    out.release()
+    return 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
@@ -282,17 +362,20 @@ 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
+        return f"Error: Failed to process video\nDebug Log:\n{debug_log}", None, 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)
+        return (f"Error: {trajectory_log}\nDebug Log:\n{debug_log}", None, None, None, None)
 
     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_2d, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path)
+    output_path_2d = f"output_2d_{uuid.uuid4()}.mp4"
+    slow_motion_path = generate_slow_motion(frames, trajectory_2d, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path_2d)
+
+    output_path_3d = f"output_3d_{uuid.uuid4()}.mp4"
+    trajectory_video_path = generate_3d_trajectory_video(trajectory_3d, pitch_point_3d, impact_point_3d, detection_frames, pitch_frame, impact_frame, output_path_3d)
 
     detections_fig = None
     trajectory_fig = None
@@ -303,6 +386,7 @@ def drs_review(video):
     debug_output = f"{debug_log}\n{trajectory_log}"
     return (f"DRS Decision: {decision}\nDebug Log:\n{debug_output}", 
             slow_motion_path, 
+            trajectory_video_path,
             detections_fig, 
             trajectory_fig)
 
@@ -312,12 +396,13 @@ 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="2D Slow-Motion Replay with Ball Detection (Green), Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow)"),
+        gr.Video(label="3D Trajectory Video with Trajectory (Blue Line), Pitch Point (Red), Impact Point (Yellow), Wicket Lines (Black)"),
         gr.Plot(label="3D Ball Detections Plot"),
         gr.Plot(label="3D Ball Trajectory Plot")
     ],
     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, a 2D slow-motion replay, a 3D trajectory video, and 3D visualizations. The 2D replay shows ball detection (green boxes), trajectory (blue line), pitch point (red circle), and impact point (yellow circle). The 3D video and plots show the trajectory (blue line), detections (green markers), pitch point (red), impact point (yellow), and wicket lines (black)."
 )
 
 if __name__ == "__main__":