diff --git a/config.yaml b/config.yaml index 8c20fd2d..f43f0d17 100644 --- a/config.yaml +++ b/config.yaml @@ -32,10 +32,11 @@ skip_process: false # if true, skips processing skip_court: false # if true, skips court detection skip_player_filter: false # if true, allows all tracked ids to be accounted -# Processing parameters +# Data cleaning parameters filter_threshold: 10 # min frames for player to be considered in possession join_threshold: 20 # max frames for same player to still be in possession shot_window: 10 # window of frames to look at ball intersecting top box and rim box +ball_window: 30 # symmetric windows of frames to consider ball # Pose action parameters shot_threshold: 0.8 # threshold for shot action diff --git a/data/trimmed.mov b/data/trimmed.mov new file mode 100644 index 00000000..c794da8e Binary files /dev/null and b/data/trimmed.mov differ diff --git a/src/processing/clean.py b/src/processing/clean.py new file mode 100644 index 00000000..0c5ca54c --- /dev/null +++ b/src/processing/clean.py @@ -0,0 +1,37 @@ +from state import GameState, Frame + + +class Clean: + def __init__(self, state: GameState): + self.state = state + + def run(self, window: int): + self.clean_ball(window) + + def incr(self, d: dict, fr: Frame, incr=1): + "increments frequency of seen ball in frame by 1" + for b in fr.ball_candidates: + d[b] = d.get(b, 0) + incr + + def decr(self, d: dict, fr: Frame): + "decrements frequency of seen ball in frame by 1" + self.incr(d, fr, incr=-1) + + def clean_ball(self, window: int): + "assigns ball with highest frame frequency over a window" + freq: dict[str, int] = {} # freq of id over winow + frames = self.state.frames + window = int(window / 2) + for i in range(min(window, len(frames))): + self.incr(freq, frames[i]) + for i, cur_fr in enumerate(frames): + if i + window < len(frames): # add new + self.incr(freq, frames[i + window]) + max_v = max(freq.values()) + max_ks = [k for k, v in freq.items() if v == max_v] # best keys + for k in max_ks: + if k in cur_fr.ball_candidates: + cur_fr.ball = cur_fr.ball_candidates[k] + break + if i - window >= 0: + self.decr(freq, frames[i - window]) # remove old diff --git a/src/processing/parse.py b/src/processing/parse.py index 5b810137..2d2a196f 100644 --- a/src/processing/parse.py +++ b/src/processing/parse.py @@ -5,6 +5,7 @@ from state import GameState, Frame, ObjectType, Box from pose_estimation.pose_estimate import AngleNames, KeyPointNames + def parse_sort_output(state: GameState, sort_output) -> None: """ Reads the SORT output and updates state.states frame-by-frame. @@ -23,24 +24,21 @@ def parse_sort_output(state: GameState, sort_output) -> None: sts = state.frames b = 0 # index of line in ball - s = 0 # index of state + s = 0 # index AND FRAME of state while b < len(lines): frame, obj_type, id, xmin, ymin, xwidth, ywidth = lines[b][:7] if s >= len(sts): # s-1 frameno < bframe, s = len(states) - sts.append(Frame(frame)) - elif frame < sts[s].frameno: # s-1 frameno < bframe < s frameno - sts.insert(s, Frame(frame)) - elif frame > sts[s].frameno: - if sts[s].rim is None and s > 0: - sts[s].rim = sts[s - 1].rim # ensure rim set + sts.append(Frame(s)) # append at index s + sts[s].rim = sts[s - 1].rim # ensure rim set + if frame > s: s += 1 continue sF: Frame = sts[s] - assert sF.frameno == frame + assert s == frame box = (xmin, ymin, xmin + xwidth, ymin + ywidth) if obj_type is ObjectType.BALL.value: - sF.set_ball_frame(id, *box) + sF.add_ball_frame(id, *box) elif obj_type is ObjectType.PLAYER.value: sF.add_player_frame(id, *box) elif obj_type is ObjectType.RIM.value: @@ -67,8 +65,8 @@ def parse_pose_output(state: GameState, pose_output: str) -> None: lines = [[int(x) for x in line.split()] for line in file.readlines()] file.close() - kpn = len(KeyPointNames.list) # number of keypoints - an = len(AngleNames.list) # number of angles + kpn = len(KeyPointNames.list) # number of keypoints + an = len(AngleNames.list) # number of angles sts = state.frames p = 0 # index of pose_data @@ -91,7 +89,7 @@ def parse_pose_output(state: GameState, pose_output: str) -> None: likely_id = [None, -1] for id, pf in state_frame.players.items(): pbox = pf.box - + # Calculate the area of intersection between the person's box and the player's box. intersection_area = bbox.area_of_intersection(pbox) # If the intersection area is greater than the current maximum, update the most likely player ID and area. @@ -101,10 +99,10 @@ def parse_pose_output(state: GameState, pose_output: str) -> None: # If a likely player is found, set the keypoints for that player. Otherwise, print a message. if likely_id[0] is not None: state_frame.players[likely_id[0]].set_keypoints( - lines[p][7:7+2*kpn] + lines[p][7 : 7 + 2 * kpn] ) state_frame.players[likely_id[0]].set_angles( - lines[p][7+2*kpn:7+2*kpn+an] + lines[p][7 + 2 * kpn : 7 + 2 * kpn + an] ) else: print("No likely player found for this person") diff --git a/src/processing/possession.py b/src/processing/possession.py index 75e4626c..779a8169 100644 --- a/src/processing/possession.py +++ b/src/processing/possession.py @@ -19,9 +19,9 @@ def compute_possessions(self): self._create_possession_intervals(min_length=15) self._filter_intervals() - for interval in self.possessions: - print( - f"Player {interval.playerid}: Start {interval.start}, End {interval.end}, Length {interval.length}") + # for interval in self.possessions: + # print( + # f"Player {interval.playerid}: Start {interval.start}, End {interval.end}, Length {interval.length}") return self.possessions diff --git a/src/processing/trendline.py b/src/processing/trendline.py index 4f2aa32a..b28ab5c5 100644 --- a/src/processing/trendline.py +++ b/src/processing/trendline.py @@ -5,10 +5,11 @@ def __init__(self, state: GameState, args): self.args = args self.state = state self.video_path = args["video_file"] - self.fps = 30 # Assume FPS is 30 + self.fps = 30 + self.velocity_smoothing = 3 def calculate_velocity(self): - # Calculate the velocity between each pair of frames + velocities = [] for i in range(1, len(self.state.frames)): frame1 = self.state.frames[i - 1] frame2 = self.state.frames[i] @@ -17,16 +18,22 @@ def calculate_velocity(self): x_center1, y_center1 = frame1.ball.box.center() x_center2, y_center2 = frame2.ball.box.center() - # Time difference in seconds (assuming fps is 30) time_diff = 1 / self.fps - - # Velocity components + time_diff = 1 vx = (x_center2 - x_center1) / time_diff vy = (y_center2 - y_center1) / time_diff - # Store these velocities in frame2's ball - frame2.ball.vx = vx - frame2.ball.vy = vy + velocities.append((vx, vy)) + + if len(velocities) > self.velocity_smoothing: + velocities.pop(0) + + avg_vx = sum(v[0] for v in velocities) / len(velocities) + avg_vy = sum(v[1] for v in velocities) / len(velocities) + # print(f"Frame {i}: Velocity - vx: {avg_vx:.2f}, vy: {avg_vy:.2f}") + + frame2.ball.vx = avg_vx + frame2.ball.vy = avg_vy def estimate_missing_positions(self): for i in range(len(self.state.frames)): @@ -74,7 +81,61 @@ def create_predicted_box(self, x_center, y_center, ball_size=20): return Box(xmin_pred, ymin_pred, xmax_pred, ymax_pred, predicted=True) + # tried using acceleration, didn't work + # def calculate_acceleration(self): + # for i in range(2, len(self.state.frames)): + # frame0 = self.state.frames[i - 2] + # frame1 = self.state.frames[i - 1] + # frame2 = self.state.frames[i] + + # if frame0.ball and frame1.ball and frame2.ball: + # vx1 = frame1.ball.vx + # vy1 = frame1.ball.vy + + # vx2 = frame2.ball.vx + # vy2 = frame2.ball.vy + + # time_diff = 1 / self.fps + # time_diff = 1 + + # ax = (vx2 - vx1) / time_diff + # ay = (vy2 - vy1) / time_diff + # # print(f"Frame {i}: Acceleration - ax: {ax:.2f}, ay: {ay:.2f}") + + # frame2.ball.ax = ax + # frame2.ball.ay = ay + + # def detect_abrupt_changes(self, acceleration_threshold=500): + # for i in range(len(self.state.frames)): + # frame = self.state.frames[i] + # if frame.ball and hasattr(frame.ball, 'ax') and hasattr(frame.ball, 'ay'): + # # Dynamic threshold + # dynamic_threshold = acceleration_threshold * (1 + (abs(frame.ball.vx) + abs(frame.ball.vy)) / 100) + # if abs(frame.ball.ax) > dynamic_threshold or abs(frame.ball.ay) > dynamic_threshold: + # frame.ball = None + + def is_spatial_change_abrupt(self, current_frame, spatial_threshold=70, window_size=15): + if current_frame.ball: + x2, y2 = current_frame.ball.box.center() + + # Look back up to 30 frames + for j in range(1, min(window_size + 1, current_frame.frameno)): + prev_frame = self.state.frames[current_frame.frameno - j] + if prev_frame.ball: + x1, y1 = prev_frame.ball.box.center() + + distance = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5 + if distance > spatial_threshold: + print(f"Frame {current_frame.frameno}: Abrupt spatial change detected. Distance: {distance:.2f}") + return True + break + + return False + def process(self): self.calculate_velocity() + for i in range(len(self.state.frames)): + if self.is_spatial_change_abrupt(self.state.frames[i]): + self.state.frames[i].ball = None self.estimate_missing_positions() - return self.state \ No newline at end of file + return self.state diff --git a/src/processrunner.py b/src/processrunner.py index 28a55ba1..04efafe1 100644 --- a/src/processrunner.py +++ b/src/processrunner.py @@ -3,7 +3,18 @@ """ import state from state import GameState -from processing import parse, court, render, shot, team, video, trendline, action, possession +from processing import ( + parse, + clean, + court, + render, + shot, + team, + video, + trendline, + action, + possession, +) from args import DARGS @@ -21,14 +32,17 @@ def __init__(self, args=DARGS): def run_parse(self): "Runs parse module over SORT (and pose later) outputs to update GameState" parse.parse_sort_output(self.state, self.args["people_file"]) + parse.parse_sort_output(self.state, self.args["ball_file"]) + parse.parse_pose_output(self.state, self.args["pose_file"]) + + def run_cleaning(self): + "runs clearning module" self.state.recompute_frame_count() if not self.args["skip_player_filter"]: # in case of short video threshold = min(300, len(self.state.frames) / 3) self.state.filter_players(threshold=threshold) - - parse.parse_sort_output(self.state, self.args["ball_file"]) - parse.parse_pose_output(self.state, self.args["pose_file"]) + clean.Clean(self.state).run(self.args["ball_window"]) def run_possession(self): """self.state.recompute_possesssions() @@ -37,7 +51,8 @@ def run_possession(self): join_threshold=self.args["join_threshold"], )""" possession_computer = possession.PossessionComputer( - self.state.frames, self.state.players) # Assuming frames is a list of frame objects + self.state.frames, self.state.players + ) # Assuming frames is a list of frame objects self.state.possessions = possession_computer.compute_possessions() self.state.recompute_pass_from_possession() @@ -63,8 +78,7 @@ def run_video_render(self, homography): return videoRender = render.VideoRender(homography) videoRender.render_video(self.state, self.args["minimap_file"]) - videoRender.reencode( - self.args["minimap_file"], self.args["minimap_temp_file"]) + videoRender.reencode(self.args["minimap_file"], self.args["minimap_temp_file"]) def run_video_processor(self): video_creator = video.VideoCreator( @@ -83,12 +97,14 @@ def run(self): """ self.run_parse() print("parsing complete!") + self.run_cleaning() + print("cleaning complete!") + self.run_trendline() + print("trendline processing complete!") self.run_possession() print("possession detection complete!") self.run_team_detect() print("team detection complete!") - self.run_trendline() - print("trendline processing complete!") self.run_shot_detect() print("shot detection complete!") self.run_courtline_detect() diff --git a/src/state.py b/src/state.py index d924105b..f731897d 100644 --- a/src/state.py +++ b/src/state.py @@ -228,13 +228,14 @@ class BallFrame: vy: velocity in the y-direction """ - def __init__(self, xmin: int, ymin: int, xmax: int, ymax: int) -> None: + def __init__( + self, xmin: int, ymin: int, xmax: int, ymax: int, id: str = None + ) -> None: # IMMUTABLE self.box: Box = Box(xmin, ymin, xmax, ymax) # Bounding box # MUTABLE - self.playerid: int = None # Last player in possession - self.type: BallType = None # IN_POSSESSION, IN_TRANSITION, or OUT_OF_PLAY + self.ballid: str = id self.vx: float = None self.vy: float = None @@ -291,7 +292,7 @@ def set_keypoints(self, keypoints: list) -> None: return for i in range(len(KeyPointNames.list)): - x, y = keypoints[2 * i: 2 * i + 2] + x, y = keypoints[2 * i : 2 * i + 2] confidence = 1 # can put into confidence later key = KeyPointNames.list[i] self.keypoints[key] = Keypoint(x, y, confidence) @@ -399,6 +400,8 @@ def __init__(self, frameno: int) -> None: self.players: dict[str, PlayerFrame] = {} "dictionary of form {player_[id] : PlayerFrame}" self.ball: BallFrame = None # ASSUMPTION: SINGLE BALLS + "ball of the frame" + self.ball_candidates: dict[str, BallFrame] = {} "dictionary of form {ball_[id] : BallFrame}" self.rim: Box = None # ASSUMPTION: SINGLE RIM "bounding box of rim" @@ -411,11 +414,11 @@ def add_player_frame(self, id: int, xmin: int, ymin: int, xmax: int, ymax: int): id = "player_" + str(id) self.players.update({id: pf}) - def set_ball_frame(self, id: int, xmin: int, ymin: int, xmax: int, ymax: int): + def add_ball_frame(self, id: int, xmin: int, ymin: int, xmax: int, ymax: int): "set ball in frame given id and bounding boxes" - bf = BallFrame(xmin, ymin, xmax, ymax) id = "ball_" + str(id) - self.ball = bf + bf = BallFrame(xmin, ymin, xmax, ymax, id) + self.ball_candidates.update({id: bf}) def set_rim_box(self, id: int, xmin: int, ymin: int, xmax: int, ymax: int): "set rim box given bounding boxes"