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
 import tempfile
@@ -13,7 +12,7 @@ import tempfile
 model = YOLO("best.pt")
 model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Dynamically resolve ball class index
+# Resolve the class index for "cricketBall"
 ball_class_index = None
 for k, v in model.names.items():
     if v.lower() == "cricketball":
@@ -34,6 +33,26 @@ PITCH_LENGTH = 20.12
 STUMPS_HEIGHT = 0.71
 MAX_POSITION_JUMP = 30
 
+def bezier_curve(points, num=100):
+    """Compute a quadratic Bezier curve from given points."""
+    points = np.array(points)
+    if len(points) < 3:
+        return points
+    p0, p1, p2 = points[0], points[len(points)//2], points[-1]
+    t = np.linspace(0, 1, num=num)
+    curve = (1 - t)[:, None]**2 * p0 + 2 * (1 - t)[:, None] * t[:, None] * p1 + t[:, None]**2 * p2
+    return curve
+
+def draw_stumps_overlay(frame):
+    height, width = frame.shape[:2]
+    stumps_x = width // 2
+    stump_top = int(height * 0.1)
+    stump_bottom = int(height * 0.9)
+    stump_width_px = int(width * 0.03)
+    for offset in [-stump_width_px, 0, stump_width_px]:
+        cv2.line(frame, (stumps_x + offset, stump_top), (stumps_x + offset, stump_bottom), (0, 0, 0), 2)
+    return frame
+
 def process_video(video_path):
     if not os.path.exists(video_path):
         return [], [], [], "Error: Video file not found"
@@ -42,7 +61,6 @@ def process_video(video_path):
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     frames, ball_positions, detection_frames, debug_log = [], [], [], []
     frame_count = 0
-
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
@@ -62,23 +80,11 @@ def process_video(video_path):
         frames[-1] = frame
         debug_log.append(f"Frame {frame_count}: {detections} ball detections")
     cap.release()
-
-    if not ball_positions:
-        debug_log.append("No balls detected in any frame")
-    else:
-        debug_log.append(f"Total ball detections: {len(ball_positions)}")
-        debug_log.append(f"Video resolution: {frame_width}x{frame_height}")
-
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
 def find_bounce_point(ball_coords):
-    """
-    Detect bounce point using y-derivative reversal with early-frame suppression.
-    Looks for where y increases then decreases (ball hits ground).
-    """
     y_coords = [p[1] for p in ball_coords]
     min_index = None
-
     for i in range(2, len(y_coords) - 2):
         dy1 = y_coords[i] - y_coords[i - 1]
         dy2 = y_coords[i + 1] - y_coords[i]
@@ -86,89 +92,70 @@ def find_bounce_point(ball_coords):
             if i > len(y_coords) * 0.2:
                 min_index = i
                 break
-
-    if min_index is not None:
-        return ball_coords[min_index]
-    
-    return ball_coords[len(ball_coords)//2]
-
-def lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point):
-    if not frames or not trajectory or len(ball_positions) < 2:
-        return "Not enough data", trajectory, pitch_point, impact_point
-
-    frame_height, frame_width = frames[0].shape[:2]
-    stumps_x = frame_width / 2
-    stumps_y = frame_height * 0.9
-    stumps_width_pixels = frame_width * (STUMPS_WIDTH / 3.0)
-
-    pitch_x, _ = pitch_point
-    impact_x, impact_y = impact_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)", 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)", 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 projected to hit stumps)", trajectory, pitch_point, impact_point
-    return f"Not Out (Missing stumps)", trajectory, pitch_point, impact_point
+    return ball_coords[min_index] if min_index else ball_coords[len(ball_coords)//2]
 
 def estimate_trajectory(ball_positions, detection_frames, frame_height, frame_width):
     if len(ball_positions) < 2:
         return None, None, None, "Error: Not enough ball detections"
-
     filtered_positions = [ball_positions[0]]
     filtered_frames = [detection_frames[0]]
     for i in range(1, len(ball_positions)):
-        prev, curr = filtered_positions[-1], ball_positions[i]
-        if np.linalg.norm(np.array(curr) - np.array(prev)) <= MAX_POSITION_JUMP:
-            filtered_positions.append(curr)
+        if np.linalg.norm(np.array(ball_positions[i]) - np.array(filtered_positions[-1])) <= MAX_POSITION_JUMP:
+            filtered_positions.append(ball_positions[i])
             filtered_frames.append(detection_frames[i])
-
     if len(filtered_positions) < 2:
         return None, None, None, "Error: Filtered detections too few"
-
     x_vals = [p[0] for p in filtered_positions]
     y_vals = [p[1] for p in filtered_positions]
     times = np.array(filtered_frames) / FRAME_RATE
-
     try:
         fx = interp1d(times, x_vals, kind='cubic', fill_value="extrapolate")
         fy = interp1d(times, y_vals, kind='cubic', fill_value="extrapolate")
     except Exception as e:
         return None, None, None, f"Interpolation error: {str(e)}"
-
-    total_frames = max(filtered_frames) - min(filtered_frames) + 1
-    t_full = np.linspace(times[0], times[-1], max(5, total_frames * SLOW_MOTION_FACTOR))
-    x_full = fx(t_full)
-    y_full = fy(t_full)
+    t_full = np.linspace(times[0], times[-1], max(5, len(times) * SLOW_MOTION_FACTOR))
+    x_full, y_full = fx(t_full), fy(t_full)
     trajectory = list(zip(x_full, y_full))
-
     pitch_point = find_bounce_point(filtered_positions)
     impact_point = filtered_positions[-1]
-
     return trajectory, pitch_point, impact_point, "Trajectory estimated successfully"
 
+def lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point):
+    if not frames or not trajectory or len(ball_positions) < 2:
+        return "Not enough data", trajectory, pitch_point, impact_point
+    frame_height, frame_width = frames[0].shape[:2]
+    stumps_x = frame_width / 2
+    stumps_y = frame_height * 0.9
+    stumps_width_pixels = frame_width * (STUMPS_WIDTH / 3.0)
+    pitch_x, _ = pitch_point
+    impact_x, _ = impact_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)", 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)", 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 projected to hit stumps)", trajectory, pitch_point, impact_point
+    return f"Not Out (Missing stumps)", trajectory, pitch_point, impact_point
+
 def generate_replay(frames, trajectory, pitch_point, impact_point, detection_frames):
     if not frames or not trajectory:
         return None
-
-    temp_file = os.path.join(tempfile.gettempdir(), f"drs_output_{uuid.uuid4()}.mp4")
+    bezier = bezier_curve(trajectory)
+    out_path = os.path.join(tempfile.gettempdir(), f"drs_output_{uuid.uuid4()}.mp4")
     height, width = frames[0].shape[:2]
-    out = cv2.VideoWriter(temp_file, cv2.VideoWriter_fourcc(*'mp4v'), FRAME_RATE / SLOW_MOTION_FACTOR, (width, height))
-
-    min_frame = min(detection_frames)
-    max_frame = max(detection_frames)
+    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), FRAME_RATE / SLOW_MOTION_FACTOR, (width, height))
+    min_frame, max_frame = min(detection_frames), max(detection_frames)
     total_frames = max_frame - min_frame + 1
-    traj_per_frame = max(1, len(trajectory) // total_frames)
-    indices = [min(i * traj_per_frame, len(trajectory)-1) for i in range(total_frames)]
-
+    traj_per_frame = max(1, len(bezier) // total_frames)
+    indices = [min(i * traj_per_frame, len(bezier)-1) for i in range(total_frames)]
     for i, frame in enumerate(frames):
+        frame = draw_stumps_overlay(frame)
         idx = i - min_frame
         if 0 <= idx < len(indices):
             end_idx = indices[idx]
-            points = np.array(trajectory[:end_idx+1], dtype=np.int32).reshape((-1, 1, 2))
-            cv2.polylines(frame, [points], False, (255, 0, 0), 2)
+            pts = np.array(bezier[:end_idx+1], dtype=np.int32).reshape((-1, 1, 2))
+            cv2.polylines(frame, [pts], False, (255, 0, 0), 2)
         if pitch_point and i == detection_frames[0]:
             cv2.circle(frame, tuple(map(int, pitch_point)), 6, (0, 0, 255), -1)
         if impact_point and i == detection_frames[-1]:
@@ -176,21 +163,18 @@ def generate_replay(frames, trajectory, pitch_point, impact_point, detection_fra
         for _ in range(SLOW_MOTION_FACTOR):
             out.write(frame)
     out.release()
-    return temp_file
+    return out_path
 
 def drs_review(video):
     frames, ball_positions, detection_frames, debug_log = process_video(video)
     if not frames or not ball_positions:
         return "No frames or detections found.", None
-
     frame_height, frame_width = frames[0].shape[:2]
     trajectory, pitch_point, impact_point, log = estimate_trajectory(ball_positions, detection_frames, frame_height, frame_width)
     if not trajectory:
         return f"{log}\n{debug_log}", None
-
     decision, _, _, _ = lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point)
     replay_path = generate_replay(frames, trajectory, pitch_point, impact_point, detection_frames)
-
     result_log = f"DRS Decision: {decision}\n\n{log}\n\n{debug_log}"
     return result_log, replay_path
 
@@ -200,10 +184,10 @@ iface = gr.Interface(
     inputs=gr.Video(label="Upload Cricket Delivery Video"),
     outputs=[
         gr.Textbox(label="DRS Result and Debug Info"),
-        gr.Video(label="Replay with Trajectory & Decision")
+        gr.Video(label="Replay with Bezier Trajectory, Pitch & Impact Points")
     ],
-    title="GullyDRS - AI-Powered LBW Review",
-    description="Upload a cricket delivery video. The system will track the ball, estimate trajectory, and return a replay with an OUT/NOT OUT decision."
+    title="GullyDRS - Enhanced DRS with Bezier Trajectory",
+    description="Upload a cricket delivery video. System overlays stumps, draws Bezier trajectory, and returns annotated replay with DRS decision."
 )
 
 if __name__ == "__main__":