Pose JSON Parsing

- Create parse_pose_output in parse.py to read the pose output JSON and update the GameState
- Create Keypoint class in state.py
- Add self.keypoints to PlayerFrame
- Add set_keypoints to PlayerFrame
- Need to test: supposed to assign keypoints based on most overlap with sort bounding boxes and yolo-pose bounding boxes

.gitignore

@@ -10,3 +10,4 @@ tmp/
 tmp/*.json
 .DS_Store
 .vscode/
+.DS_Store

src/processing/parse.py

@@ -2,8 +2,8 @@
 Parsing module for parsing all
 models outputs into the state
 """
-from state import GameState, Frame, ObjectType
-
+from state import GameState, Frame, ObjectType, Box
+import json
 
 def parse_sort_output(state: GameState, sort_output) -> None:
     """
@@ -47,3 +47,64 @@ def parse_sort_output(state: GameState, sort_output) -> None:
             sF.set_rim_box(id, *box)
 
         b += 1  # process next line
+
+def parse_pose_output(state: GameState, pose_data_json: str) -> None:
+    """
+    Parses pose data and updates the player frames in the game state.
+
+    Inputs:
+    - state (GameState): The game state object to be updated.
+    - pose_data_json (str): File path to the pose data JSON file.
+
+    Notes:
+    - Pose data is expected to be in JSON format with keypoint information for each player in each frame.
+    - The function assumes the frames in the pose data are in the same order as in the game state.
+    """
+    # Open the pose data JSON file and load its content.
+    with open(pose_data_json, "r") as f:
+        pose_data = json.load(f)
+
+    # Iterate through each frame's data in the pose data.
+    for frame_no, frame_data in enumerate(pose_data):
+        try:
+            # Try to retrieve the corresponding frame from the game state using the frame number.
+            frame = state.frames[frame_no]
+        except IndexError:
+            # If the frame number is not found in the game state, print an error message and continue to the next iteration.
+            print(f"Frame number {frame_no} not found in the game state.")
+            continue
+
+        # Retrieve the list of persons (players) from the current frame's data. If not found, default to an empty list.
+        frame_persons = frame_data.get('persons', [])
+        if not frame_persons:
+            # If there are no persons found in the frame, print a message and continue to the next iteration.
+            print("No persons data found in this frame")
+            continue
+
+        # Iterate through each person in the frame.
+        for person in frame_persons:
+            # Try to retrieve the bounding box of the person. If not found, print an error message and continue to the next iteration.
+            bounding_box = person.get('box')
+            if bounding_box is None:
+                print("No bounding box data found for this person")
+                continue
+
+            # Unpack the bounding box coordinates and create a Box object.
+            xmin, ymin, xmax, ymax = bounding_box
+            pose_box = Box(xmin, ymin, xmax, ymax)
+
+            # Initialize variables to keep track of the most likely player ID and the area of intersection.
+            likely_id = [None, -1]
+            # Iterate through each player in the frame.
+            for player_id, player_frame in frame.players.items():
+                # Calculate the area of intersection between the person's box and the player's box.
+                intersection_area = pose_box.area_of_intersection(player_frame.box)
+                # If the intersection area is greater than the current maximum, update the most likely player ID and area.
+                if intersection_area > likely_id[1]:
+                    likely_id = [player_id, intersection_area]
+
+            # If a likely player is found, set the keypoints for that player. Otherwise, print a message.
+            if likely_id[0] is not None:
+                frame.players[likely_id[0]].set_keypoints(person.get('keypoints'))
+            else:
+                print("No likely player found for this person")

src/state.py

@@ -213,20 +213,55 @@ def __init__(self, xmin: int, ymin: int, xmax: int, ymax: int) -> None:
         "ball in possession (-1) if not in possession"
         self.type: ActionType = None
         "NOTHING, DRIBBLE, PASS, SHOOT"
+        self.keypoints: dict[str, Keypoint] = {}
+        "keypoints of the player"
+
+    def set_keypoints(self, keypoints_data: dict) -> None:
+        "Sets the keypoints for the player"
+        for key, value in keypoints_data.items():
+            x, y, confidence = value
+            self.keypoints[key] = Keypoint(x, y, confidence)
 
     def check(self) -> bool:
         "verifies if well-defined"
         try:
-            assert self.box.check() == True
+            assert self.box.check()
             assert self.type is not None
             if self.type is ActionType.NOTHING:
                 assert self.ballid == -1
             else:
                 assert self.ballid != -1
-        except:
+            for keypoint in self.keypoints.values():
+                assert keypoint.check()
+        except AssertionError:
             return False
         return True
 
+class Keypoint:
+    """
+    Keypoint class containing the coordinates and confidence of a keypoint
+    """
+
+    def __init__(self, x: float, y: float, confidence: float) -> None:
+        # IMMUTABLE
+        self.x: float = x
+        "x-coordinate of the keypoint"
+        self.y: float = y
+        "y-coordinate of the keypoint"
+        self.confidence: float = confidence
+        "confidence score of the keypoint detection"
+
+    def __repr__(self) -> str:
+        return f"Keypoint(x={self.x}, y={self.y}, confidence={self.confidence})"
+
+    def check(self) -> bool:
+        "verifies if well-defined"
+        try:
+            assert 0 <= self.confidence <= 1
+            assert self.x >= 0 and self.y >= 0
+        except AssertionError:
+            return False
+        return True
 
 class Frame:
     "Frame class containing frame-by-frame information"

src/view/app.py

@@ -20,7 +20,7 @@
 if "state" not in st.session_state:
     st.session_state.state = 0
     st.session_state.logo = "src/view/static/basketball.png"
-    with open("data/short_new_1.mp4", "rb") as file:
+    with open("data/training_data.mp4", "rb") as file:
         st.session_state.video_file = io.BytesIO(file.read())
     st.session_state.processed_video = None
     st.session_state.result_string = None