Update gully_drs_core/ball_detection.py

gully_drs_core/ball_detection.py

@@ -5,14 +5,23 @@ import numpy as np
 from .model_utils import load_model
 
 def find_bounce_point(path):
-    """
-    Finds the point where the ball bounces by detecting a Y-axis dip.
-    """
     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):
+    if len(ball_path) < 2:
+        return 0.0
+    p1 = ball_path[0]
+    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
+    speed_kmh = (dist_m / time_s) * 3.6 if time_s > 0 else 0
+    return round(speed_kmh, 1)
+
 def analyze_video(file_path):
     model = load_model()
     cap = cv2.VideoCapture(file_path)
@@ -23,44 +32,67 @@ def analyze_video(file_path):
     ball_path = []
     frames = []
 
+    max_jump = 100  # max allowed jump (pixels) between consecutive ball detections
+
+    last_point = None
+    frame_idx = 0
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
 
         results = model(frame)
+        valid_detection = None
+
         for r in results:
-            for box in r.boxes:
-                cls = int(box.cls[0])
-                if cls == 32:  # cricket ball class
-                    x1, y1, x2, y2 = map(int, box.xyxy[0])
-                    cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
-                    ball_path.append((cx, cy))
-                    cv2.circle(frame, (cx, cy), 6, (0, 255, 0), -1)
+            # 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, cy = (x1 + x2) // 2, (y1 + y2) // 2
+
+                # Check jump threshold from last point
+                if last_point:
+                    dx, dy = cx - last_point[0], cy - last_point[1]
+                    jump = (dx**2 + dy**2)**0.5
+                    if jump > max_jump:
+                        # Reject outlier
+                        frames.append(frame)
+                        frame_idx += 1
+                        continue
+
+                valid_detection = (cx, cy)
+                last_point = valid_detection
+                cv2.circle(frame, valid_detection, 6, (0, 255, 0), -1)
+
+        if valid_detection:
+            ball_path.append(valid_detection)
 
         frames.append(frame)
+        frame_idx += 1
 
     cap.release()
 
-    # Analyze trajectory
     bounce_point = find_bounce_point(ball_path)
     impact_point = ball_path[-1] if ball_path else None
 
-    # Stump zone: middle area at bottom
     stump_zone = (
-        width // 2 - 30,  # x1
-        height - 100,     # y1
-        width // 2 + 30,  # x2
-        height            # y2
+        width // 2 - 30,
+        height - 100,
+        width // 2 + 30,
+        height
     )
 
-    # LBW Decision: if impact is 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,
@@ -68,5 +100,6 @@ def analyze_video(file_path):
         "bounce_point": bounce_point,
         "impact_point": impact_point,
         "decision": decision,
-        "stump_zone": stump_zone
+        "stump_zone": stump_zone,
+        "speed_kmh": speed_kmh
     }