GullyDRS / gully_drs_core /ball_detection.py
AjaykumarPilla's picture
Update gully_drs_core/ball_detection.py
0ae5282 verified
# gully_drs_core/ball_detection.py
import cv2
import numpy as np
from .model_utils import load_model
def find_bounce_point(path):
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)
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
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:
# 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()
bounce_point = find_bounce_point(ball_path)
impact_point = ball_path[-1] if ball_path else None
stump_zone = (
width // 2 - 30,
height - 100,
width // 2 + 30,
height
)
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,
"frames": frames,
"bounce_point": bounce_point,
"impact_point": impact_point,
"decision": decision,
"stump_zone": stump_zone,
"speed_kmh": speed_kmh
}