File size: 1,757 Bytes
b575849 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# Data Classes (Don't pass dicts everywhere!)
from dataclasses import dataclass
from typing import List, Optional, Tuple
from enum import Enum
class EventType(Enum):
PASS = "pass"
DRIBBLE = "dribble"
MOVEMENT = "movement"
RECOVERY = "recovery"
SHOT = "shot"
@dataclass
class Detection:
"""Single detection from RF-DETR"""
class_id: int
confidence: float
bbox: Tuple[float, float, float, float] # x, y, width, height
class_name: str = ""
@dataclass
class TrackedObject:
"""Tracked object with ID"""
object_id: int
detection: Detection
team_id: Optional[int] = None
@dataclass
class Player:
"""Player with pitch coordinates"""
object_id: int
team_id: int
x_pitch: float
y_pitch: float
bbox: Tuple[float, float, float, float]
frame_id: int
timestamp: float
@dataclass
class Ball:
"""Ball position"""
x_pitch: float
y_pitch: float
bbox: Tuple[float, float, float, float]
frame_id: int
timestamp: float
object_id: Optional[int] = None
@dataclass
class FrameData:
"""Data for a single frame"""
frame_id: int
timestamp: float
players: List[Player]
ball: Optional[Ball] = None
detections: List[Detection] = None
@dataclass
class Location:
"""Pitch location"""
x: float
y: float
@dataclass
class Event:
"""Detected event"""
id: str
type: EventType
start_frame: int
end_frame: int
start_location: Location
end_location: Location
involved_players: List[int] # object_ids
confidence: float
timestamp_start: float
timestamp_end: float
@dataclass
class MatchData:
"""Complete match data"""
match_id: str
events: List[Event]
metadata: dict
|