Update app.py

app.py

@@ -1,10 +1,10 @@
-
 import cv2
 import numpy as np
 import torch
 from ultralytics import YOLO
 import gradio as gr
-from scipy.interpolate import interp1d, UnivariateSpline
+from scipy.interpolate import interp1d
+from scipy.ndimage import uniform_filter1d
 import uuid
 import os
 
@@ -15,10 +15,10 @@ model = YOLO("best.pt")
 STUMPS_WIDTH = 0.2286  # meters (width of stumps)
 FRAME_RATE = 20  # Input video frame rate
 SLOW_MOTION_FACTOR = 2  # Reduced for faster output
-CONF_THRESHOLD = 0.25  # Confidence threshold for detection
-PITCH_ZONE_Y = 0.85  # Adjusted for pitch near stumps
-IMPACT_ZONE_Y = 0.75  # Adjusted for impact near batsman leg
-IMPACT_DELTA_Y = 30  # Reduced for finer impact detection
+CONF_THRESHOLD = 0.3  # Increased for better detection
+PITCH_ZONE_Y = 0.8  # Adjusted for pitch near stumps
+IMPACT_ZONE_Y = 0.7  # Adjusted for impact near batsman leg
+IMPACT_DELTA_Y = 20  # Reduced for finer impact detection
 STUMPS_HEIGHT = 0.711  # meters (height of stumps)
 
 def process_video(video_path):
@@ -35,16 +35,18 @@ def process_video(video_path):
         ret, frame = cap.read()
         if not ret:
             break
-        if frame_count % 2 == 0:  # Process every 2nd frame
-            frames.append(frame.copy())
-            results = model.predict(frame, conf=CONF_THRESHOLD)
-            detections = [det for det in results[0].boxes if det.cls == 0]
-            if len(detections) == 1:
-                x1, y1, x2, y2 = detections[0].xyxy[0].cpu().numpy()
-                ball_positions.append([(x1 + x2) / 2, (y1 + y2) / 2])
-                detection_frames.append(len(frames) - 1)
-                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
-            frames[-1] = frame
+        # Process every frame for better tracking
+        frames.append(frame.copy())
+        # Preprocess frame for better detection
+        frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=10)  # Enhance contrast
+        results = model.predict(frame, conf=CONF_THRESHOLD)
+        detections = [det for det in results[0].boxes if det.cls == 0]
+        if len(detections) == 1:
+            x1, y1, x2, y2 = detections[0].xyxy[0].cpu().numpy()
+            ball_positions.append([(x1 + x2) / 2, (y1 + y2) / 2])
+            detection_frames.append(len(frames) - 1)
+            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}: {len(detections)} ball detections")
         frame_count += 1
     cap.release()
@@ -61,25 +63,18 @@ def estimate_trajectory(ball_positions, detection_frames, frames):
         return None, None, None, None, None, None, "Error: Fewer than 2 valid single-ball detections for trajectory"
     frame_height = frames[0].shape[0]
 
-    # Filter to unique positions
-    unique_positions = [ball_positions[0]]
-    for pos in ball_positions[1:]:
-        if abs(pos[0] - unique_positions[-1][0]) > 10 or abs(pos[1] - unique_positions[-1][1]) > 10:
-            unique_positions.append(pos)
-    x_coords = [pos[0] for pos in unique_positions]
-    y_coords = [pos[1] for pos in unique_positions]
-    times = np.array([i / FRAME_RATE for i in range(len(unique_positions))])
-
-    # Smooth coordinates with spline interpolation
-    x_smooth = UnivariateSpline(times, x_coords, s=10)
-    y_smooth = UnivariateSpline(times, y_coords, s=10)
+    # Smooth coordinates with moving average
+    window_size = 3
+    x_coords = uniform_filter1d([pos[0] for pos in ball_positions], size=window_size, mode='nearest')
+    y_coords = uniform_filter1d([pos[1] for pos in ball_positions], size=window_size, mode='nearest')
+    times = np.array([i / FRAME_RATE for i in range(len(ball_positions))])
 
     pitch_idx = 0
     for i, y in enumerate(y_coords):
         if y > frame_height * PITCH_ZONE_Y:
             pitch_idx = i
             break
-    pitch_point = unique_positions[pitch_idx]
+    pitch_point = ball_positions[pitch_idx]
     pitch_frame = detection_frames[pitch_idx]
 
     impact_idx = None
@@ -90,7 +85,7 @@ def estimate_trajectory(ball_positions, detection_frames, frames):
             break
     if impact_idx is None:
         impact_idx = len(y_coords) - 1
-    impact_point = unique_positions[impact_idx]
+    impact_point = ball_positions[impact_idx]
     impact_frame = detection_frames[impact_idx]
 
     x_coords = x_coords[:impact_idx + 1]
@@ -98,12 +93,12 @@ def estimate_trajectory(ball_positions, detection_frames, frames):
     times = times[:impact_idx + 1]
 
     try:
-        fx = interp1d(times, x_smooth(times), kind='linear', fill_value="extrapolate")
-        fy = interp1d(times, y_smooth(times), kind='quadratic', fill_value="extrapolate")
+        fx = interp1d(times, x_coords, kind='linear', fill_value="extrapolate")
+        fy = interp1d(times, y_coords, kind='quadratic', fill_value="extrapolate")
     except Exception as e:
         return None, None, None, None, None, None, f"Error in trajectory interpolation: {str(e)}"
 
-    vis_trajectory = list(zip(x_smooth(times), y_smooth(times)))
+    vis_trajectory = list(zip(x_coords, y_coords))
     t_full = np.linspace(times[0], times[-1] + 0.5, len(times) + 5)
     x_full = fx(t_full)
     y_full = fy(t_full)
@@ -122,9 +117,9 @@ def lbw_decision(ball_positions, full_trajectory, frames, pitch_point, impact_po
 
     frame_height, frame_width = frames[0].shape[:2]
     stumps_x = frame_width / 2
-    stumps_y = frame_height * 0.85  # Adjusted to align with pitch
+    stumps_y = frame_height * 0.8  # Adjusted to align with pitch
     stumps_width_pixels = frame_width * (STUMPS_WIDTH / 3.0)
-    batsman_area_y = frame_height * 0.75
+    batsman_area_y = frame_height * 0.7
 
     pitch_x, pitch_y = pitch_point
     impact_x, impact_y = impact_point
@@ -154,7 +149,7 @@ def generate_slow_motion(frames, vis_trajectory, pitch_point, pitch_frame, impac
         return None
     frame_height, frame_width = frames[0].shape[:2]
     stumps_x = frame_width / 2
-    stumps_y = frame_height * 0.85  # Align with pitch
+    stumps_y = frame_height * 0.8  # Align with pitch
     stumps_width_pixels = frame_width * (STUMPS_WIDTH / 3.0)
     stumps_height_pixels = frame_height * (STUMPS_HEIGHT / 3.0)
 
@@ -228,4 +223,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()