Ball box clean

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

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

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")

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
 

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
+        return self.state

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()

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"