Spaces:
Sleeping
Sleeping
File size: 2,378 Bytes
3bf3b3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
from ultralytics import YOLO
import cv2
import numpy as np
import os
# Load YOLO model (custom-trained or pretrained with compatible classes)
model = YOLO("yolov8n.pt") # Replace with "lbw_yolov8.pt" if custom-trained
# Target class IDs — update based on your custom model class mapping
CLASS_NAMES = {
0: "ball",
1: "bat",
2: "pad",
3: "stump",
4: "player"
}
def detect_lbw_event(frames):
"""
Detects ball, bat, stump, and pad in each frame.
Identifies impact and prepares coordinates for trajectory modeling.
Returns:
dict: {
"ball_positions": [x, y] list per frame,
"impact_frame": int,
"impact_type": str,
"objects_per_frame": [
{"ball": (x, y), "pad": (x, y), ...}
]
}
"""
ball_positions = []
impact_frame = -1
impact_type = None
objects_per_frame = []
for idx, frame in enumerate(frames):
results = model(frame)[0]
frame_objects = {}
for det in results.boxes.data:
x1, y1, x2, y2, conf, cls = det.cpu().numpy()
class_id = int(cls)
class_name = CLASS_NAMES.get(class_id, "unknown")
center_x = int((x1 + x2) / 2)
center_y = int((y1 + y2) / 2)
frame_objects[class_name] = (center_x, center_y)
if class_name == "ball":
ball_positions.append((idx, center_x, center_y))
objects_per_frame.append(frame_objects)
# Basic impact logic: ball overlaps pad or bat
if "ball" in frame_objects and ("pad" in frame_objects or "bat" in frame_objects):
bx, by = frame_objects["ball"]
if "pad" in frame_objects:
px, py = frame_objects["pad"]
if abs(bx - px) < 30 and abs(by - py) < 30:
impact_frame = idx
impact_type = "pad"
break
if "bat" in frame_objects:
tx, ty = frame_objects["bat"]
if abs(bx - tx) < 30 and abs(by - ty) < 30:
impact_frame = idx
impact_type = "bat"
break
return {
"ball_positions": ball_positions,
"impact_frame": impact_frame,
"impact_type": impact_type,
"objects_per_frame": objects_per_frame
}
|