Update gully_drs_core/ball_detection.py

gully_drs_core/ball_detection.py

@@ -5,53 +5,37 @@ import numpy as np
 from .model_utils import load_model
 
 def find_bounce_point(path):
-    """
-    Detects the bounce point by checking for a dip in the y-axis (trajectory).
-    """
-    for i in range(1, len(path) - 1):
-        if path[i - 1][1] > path[i][1] < path[i + 1][1]:  # y decreases then increases
+    for i in range(1, len(path)-1):
+        if path[i-1][1] > path[i][1] < path[i+1][1]:  # y dips = bounce
             return path[i]
     return None
 
 def estimate_speed(ball_path, fps, px_to_m=0.01):
-    """
-    Estimate speed in km/h based on pixel distance and frame rate.
-    Assumes 1 pixel ≈ 1cm (adjust px_to_m for better accuracy).
-    """
     if len(ball_path) < 2:
         return 0.0
-
     p1 = ball_path[0]
-    p2 = ball_path[min(5, len(ball_path) - 1)]  # use the 5th frame ahead
-
-    dx = p2[0] - p1[0]
-    dy = p2[1] - p1[1]
+    p2 = ball_path[min(5, len(ball_path)-1)]
+    dx, dy = p2[0] - p1[0], p2[1] - p1[1]
     dist_px = (dx**2 + dy**2)**0.5
     dist_m = dist_px * px_to_m
-    time_s = (min(5, len(ball_path) - 1)) / fps
-
+    time_s = (min(5, len(ball_path)-1)) / fps
     speed_kmh = (dist_m / time_s) * 3.6 if time_s > 0 else 0
     return round(speed_kmh, 1)
 
 def analyze_video(file_path):
-    """
-    Main processing function:
-    - Detects the ball using YOLOv8
-    - Builds trajectory from valid frames
-    - Detects bounce, impact, stump zone intersection
-    - Returns decision + video frame overlays
-    """
     model = load_model()
     cap = cv2.VideoCapture(file_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-    frames = []
     ball_path = []
+    frames = []
+
+    max_jump = 100  # max allowed jump (pixels) between consecutive ball detections
 
-    max_jump = 100  # pixels
     last_point = None
+    frame_idx = 0
 
     while True:
         ret, frame = cap.read()
@@ -62,20 +46,22 @@ def analyze_video(file_path):
         valid_detection = None
 
         for r in results:
-            ball_detections = [box for box in r.boxes if int(box.cls[0]) == 0]  # class 0 = cricket ball
+            # Accept only if exactly one detection of cricket ball class (e.g., class 0)
+            ball_detections = [box for box in r.boxes if int(box.cls[0]) == 0]
             if len(ball_detections) == 1:
                 box = ball_detections[0]
                 x1, y1, x2, y2 = map(int, box.xyxy[0])
-                cx = (x1 + x2) // 2
-                cy = (y1 + y2) // 2
+                cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
 
-                # Filter out sudden jumps in position
+                # Check jump threshold from last point
                 if last_point:
-                    dx = cx - last_point[0]
-                    dy = cy - last_point[1]
+                    dx, dy = cx - last_point[0], cy - last_point[1]
                     jump = (dx**2 + dy**2)**0.5
                     if jump > max_jump:
-                        break  # skip this frame
+                        # Reject outlier
+                        frames.append(frame)
+                        frame_idx += 1
+                        continue
 
                 valid_detection = (cx, cy)
                 last_point = valid_detection
@@ -85,15 +71,13 @@ def analyze_video(file_path):
             ball_path.append(valid_detection)
 
         frames.append(frame)
+        frame_idx += 1
 
     cap.release()
 
-    # Calculate analysis outputs
     bounce_point = find_bounce_point(ball_path)
     impact_point = ball_path[-1] if ball_path else None
-    speed_kmh = estimate_speed(ball_path, fps)
 
-    # Define stump zone area
     stump_zone = (
         width // 2 - 30,
         height - 100,
@@ -101,13 +85,14 @@ def analyze_video(file_path):
         height
     )
 
-    # LBW decision: does ball impact land in stump zone?
     decision = "OUT" if (
         impact_point and
         stump_zone[0] <= impact_point[0] <= stump_zone[2] and
         stump_zone[1] <= impact_point[1] <= stump_zone[3]
     ) else "NOT OUT"
 
+    speed_kmh = estimate_speed(ball_path, fps)
+
     return {
         "trajectory": ball_path,
         "fps": fps,