Update app.py

app.py

@@ -4,16 +4,15 @@ 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
 
-# Load YOLOv8 model and resolve class index
+# Load YOLOv8 model
 model = YOLO("best.pt")
 model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Dynamically resolve ball class index
+# Resolve class index for cricket ball
 ball_class_index = None
 for k, v in model.names.items():
     if v.lower() == "cricketball":
@@ -26,7 +25,7 @@ if ball_class_index is None:
 STUMPS_WIDTH = 0.2286
 BALL_DIAMETER = 0.073
 FRAME_RATE = 20
-SLOW_MOTION_FACTOR = 3
+SLOW_MOTION_FACTOR = 2  # Set to 1 for normal speed replay
 CONF_THRESHOLD = 0.2
 IMPACT_ZONE_Y = 0.85
 IMPACT_DELTA_Y = 50
@@ -63,55 +62,17 @@ def process_video(video_path):
         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]
         if dy1 > 0 and dy2 < 0:
             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[i]
+    return ball_coords[len(ball_coords) // 2]
 
 def estimate_trajectory(ball_positions, detection_frames, frame_height, frame_width):
     if len(ball_positions) < 2:
@@ -149,50 +110,79 @@ def estimate_trajectory(ball_positions, detection_frames, frame_height, frame_wi
 
     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_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
+
 def generate_replay(frames, trajectory, pitch_point, impact_point, detection_frames):
     if not frames or not trajectory:
-        return None
+        return None, None
 
-    temp_file = 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))
+    slow_path = os.path.join(tempfile.gettempdir(), f"drs_slow_{uuid.uuid4()}.mp4")
+    normal_path = os.path.join(tempfile.gettempdir(), f"drs_normal_{uuid.uuid4()}.mp4")
+
+    slow_writer = cv2.VideoWriter(slow_path, cv2.VideoWriter_fourcc(*'mp4v'), FRAME_RATE / SLOW_MOTION_FACTOR, (width, height))
+    normal_writer = cv2.VideoWriter(normal_path, cv2.VideoWriter_fourcc(*'mp4v'), FRAME_RATE, (width, height))
 
     min_frame = min(detection_frames)
     max_frame = 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)]
+    indices = [min(i * traj_per_frame, len(trajectory) - 1) for i in range(total_frames)]
 
     for i, frame in enumerate(frames):
+        frame_copy = frame.copy()
         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))
+            points = np.array(trajectory[:end_idx + 1], dtype=np.int32).reshape((-1, 1, 2))
             cv2.polylines(frame, [points], False, (255, 0, 0), 2)
+            cv2.polylines(frame_copy, [points], 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]:
             cv2.circle(frame, tuple(map(int, impact_point)), 6, (0, 255, 255), -1)
         for _ in range(SLOW_MOTION_FACTOR):
-            out.write(frame)
-    out.release()
-    return temp_file
+            slow_writer.write(frame)
+        normal_writer.write(frame_copy)
+
+    slow_writer.release()
+    normal_writer.release()
+    return slow_path, normal_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
+        return "No frames or detections found.", None, 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
+        return f"{log}\n{debug_log}", None, None
 
     decision, _, _, _ = lbw_decision(ball_positions, trajectory, frames, pitch_point, impact_point)
-    replay_path = generate_replay(frames, trajectory, pitch_point, impact_point, detection_frames)
+    slow_path, normal_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
+    return result_log, slow_path, normal_path
 
 # Gradio Interface
 iface = gr.Interface(
@@ -200,10 +190,11 @@ 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="Slow-Motion Replay"),
+        gr.Video(label="Normal-Speed Trajectory Only")
     ],
     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."
+    description="Upload a cricket delivery video. The system will track the ball, estimate trajectory, and return both slow-motion and normal-speed replays."
 )
 
 if __name__ == "__main__":