Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import torch
 from ultralytics import YOLO
 import gradio as gr
-from scipy.interpolate import interp1d
+from scipy.interpolate import interp1d, CubicSpline
 import uuid
 import os
 
@@ -54,31 +54,29 @@ def process_video(video_path):
 
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
-def estimate_trajectory(ball_positions, frames):
+def smooth_trajectory(ball_positions, frames):
     if len(ball_positions) < 2:
-        return None, None, None, "Error: Fewer than 2 ball detections for trajectory"
-    
-    frame_height = frames[0].shape[0]
+        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 positions up to impact for interpolation
+    # Use cubic spline interpolation to smooth the trajectory
     try:
-        fx = interp1d(times, x_coords, kind='linear', fill_value="extrapolate")
-        fy = interp1d(times, y_coords, kind='quadratic', fill_value="extrapolate")
+        spline_x = CubicSpline(times, x_coords, bc_type='natural')
+        spline_y = CubicSpline(times, y_coords, bc_type='natural')
     except Exception as e:
-        return None, None, None, f"Error in trajectory interpolation: {str(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 = fx(t_full)
-    y_full = fy(t_full)
+    x_full = spline_x(t_full)
+    y_full = spline_y(t_full)
     trajectory = list(zip(x_full, y_full))
 
-    return trajectory, None, None, "Trajectory estimated successfully"
+    return trajectory, "Trajectory smoothed successfully"
 
 def detect_pitch_and_impact(ball_positions, frames, frame_height):
     pitch_point = None
@@ -185,8 +183,9 @@ 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, pitch_point, impact_point, trajectory_log = estimate_trajectory(ball_positions, frames)
     
+    trajectory, smoothing_log = smooth_trajectory(ball_positions, frames)
+
     # Detect pitch and impact points based on ball positions
     pitch_point, impact_point = detect_pitch_and_impact(ball_positions, frames, frames[0].shape[0])