Pose estimate

.gitignore

@@ -1,3 +1,10 @@
 venv
 .env
 __pycache__
+yolov8m-pose.pt
+*.mp4
+ball/lib/python3.11/site-packages/torch/lib/libtorch_cpu.dylib
+venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu.dylib
+ball/
+tmp/
+tmp/*.json

requirements.txt

@@ -25,8 +25,8 @@ opencv-python==4.7.0.72
 matplotlib>=3.2.2
 Pillow>=7.1.2
 PyYAML>=5.3.1
-torch==2.0.1
-torchvision==0.15.2
+torch==2.0.1  # Check if this version exists, if not, use the latest stable version
+torchvision==0.15.2  # Same check as torch
 tqdm>=4.41.0
 seaborn
 scipy
@@ -46,7 +46,13 @@ imageio
 
 # View
 streamlit>=1.18.1
-hydralit_components>= 1.0.10
+hydralit_components>=1.0.10
 
 # Misc
-pylint
+pylint
+
+# Additional Dependencies for Pose Estimation
+json5
+ultralytics
+imageio==2.9.0
+imageio-ffmpeg>=0.4.3

src/main.py

@@ -36,11 +36,12 @@ def main(video_path):
 
     modelrunner = ModelRunner(video_path, model_vars)
     modelrunner.run()
-    people_output, ball_output = modelrunner.fetch_output()
+    modelrunner.pose()
+    people_output, ball_output, pose_output = modelrunner.fetch_output()
     output_video_path = 'tmp/court_video.mp4'
     output_video_path_reenc = 'tmp/court_video_reenc.mp4'
 
-    processrunner = ProcessRunner(video_path, people_output, ball_output, output_video_path, 
+    processrunner = ProcessRunner(video_path, people_output, ball_output, output_video_path,
                                   output_video_path_reenc)
     processrunner.run()
     results = processrunner.get_results()

src/modelrunner.py

@@ -6,6 +6,8 @@
 import pickle
 import subprocess
 from typing import Tuple
+from pose_estimation.pose_estimate import PoseEstimator
+from ultralytics import YOLO
 
 class ModelRunner:
     """
@@ -15,6 +17,7 @@ class ModelRunner:
     def __init__(self, video_path, model_vars) -> None:
         self.video_path = video_path
         self.frame_reduction_factor = model_vars['frame_reduction_factor']
+        self.pose_estimator = PoseEstimator(video_path=video_path)
 
 
     def drop_frames(self, input_path) -> str:
@@ -51,15 +54,24 @@ def run(self):
         with open('tmp/output.pickle', 'rb') as f:
             self.output_dict = pickle.load(f)
 
-
-    def fetch_output(self) -> Tuple[str, str]:
+    def pose(self):
+        model = YOLO('best.pt')
+        results = model(
+            source = self.video_path,
+            show=False,
+            conf=0.3,
+            verbose = False
+        )
+        self.pose_estimator.estimate_pose(results = results)
+
+    def fetch_output(self) -> Tuple[str, str, str]:
         """
         Converts the people and ball model output in self.output.dict into txt files.
         Returns a tuple of the people and ball txt output paths.
         """
-        ball_list = [tuple(round(num) for num in tup) 
+        ball_list = [tuple(round(num) for num in tup)
                      for tup in self.output_dict['basketball_data'][0]]
-        people_list = [tuple(round(num) for num in tup) 
+        people_list = [tuple(round(num) for num in tup)
                        for tup in self.output_dict['person_data'][0]]
         ball_data = [(' '.join(map(str, ball[0:7])) + ' -1 -1 -1 -1')
                      for ball in ball_list]
@@ -72,4 +84,4 @@ def fetch_output(self) -> Tuple[str, str]:
         with open('tmp/people.txt', 'w') as f:
             f.write('\n'.join(people_data))
 
-        return 'tmp/people.txt', 'tmp/ball.txt'
+        return 'tmp/people.txt', 'tmp/ball.txt', 'tmp/pose.txt'

src/pose_estimation/pose_estimate.py

@@ -0,0 +1,80 @@
+import torch
+import math
+import json
+from ultralytics import YOLO
+
+class PoseEstimator:
+    def __init__(self, model_path='best.pt', video_path='res/pose_results/test_multiple_people.mp4', combinations=None):
+        # Initialize paths, model, and combinations of keypoints to calculate angles
+        self.model_path = model_path
+        self.video_path = video_path
+        self.model = YOLO(model_path)  # Load the YOLO model
+
+        # Combinations of points to calculate 8 angles
+        self.combinations = combinations if combinations is not None else [
+            (5, 7, 9), (6, 8, 10), (11, 13, 15), (12, 14, 16),
+            (5, 6, 8), (6, 5, 7), (11, 12, 14), (12, 11, 13)
+        ]
+
+        # Names corresponding to the adjusted 8 angle types
+        self.angle_names = [
+            "left_elbow", "right_elbow", "left_knee", "right_knee",
+            "right_shoulder", "left_shoulder",
+            "right_hip", "left_hip"
+        ]
+
+    @staticmethod
+    def compute_angle(p1, p2, p3):
+        # Calculate angle given 3 points using the dot product and arc cosine
+        vector_a = p1 - p2
+        vector_b = p3 - p2
+
+        # Normalize the vectors (to make them unit vectors)
+        vector_a = vector_a / torch.norm(vector_a)
+        vector_b = vector_b / torch.norm(vector_b)
+
+        # Compute the angle
+        cosine_angle = torch.sum(vector_a * vector_b)
+        angle_radians = torch.acos(cosine_angle)
+        angle_degrees = angle_radians * 180 / math.pi
+
+        return angle_degrees
+
+    def estimate_pose(self, results):
+        model = YOLO(self.model_path)
+
+        # Initialize an empty list to store pose data
+        pose_data = []
+
+        for frame_idx, result in enumerate(results):
+            keypoints = result.keypoints.data[:, :, :2].numpy()  # Extracting the (x, y) coordinates
+            confidences = result.keypoints.conf.numpy().tolist()  # Extracting the confidences
+            boxes = result.boxes.xyxy.numpy().tolist()  # Extracting bounding boxes
+            frame_pose_data = {
+                'frame': frame_idx,
+                'persons': [],
+                'boxes': boxes,
+                'keypoints': keypoints.tolist(),
+                'confidences': confidences
+            }
+
+            for person_idx, (person_keypoints, person_confidences, box) in enumerate(zip(keypoints, confidences, boxes)):
+                person_data = {
+                    'keypoints': person_keypoints.tolist(),
+                    'confidences': person_confidences,
+                    'box': box,
+                    'angles': {}
+                }
+
+                for idx, combination in enumerate(self.combinations):
+                    if all(idx < len(person_keypoints) for idx in combination):
+                        p1, p2, p3 = (person_keypoints[i] for i in combination)
+                        angle_degrees = self.compute_angle(torch.tensor(p1), torch.tensor(p2), torch.tensor(p3))
+                        person_data['angles'][self.angle_names[idx]] = angle_degrees.item()
+
+                frame_pose_data['persons'].append(person_data)
+
+            pose_data.append(frame_pose_data)
+
+        with open("tmp/pose_data.json", "w") as f:
+            json.dump(pose_data, f)

src/processing/courtline_detect.py

@@ -318,7 +318,11 @@ def _evaluate_homography(self,pts_src:list,pts_dst:list):
         assert(pts_src is not None)
         mapped_edge_img = self._apply_gray_homography(self._MASK_COURT_EDGES,pts_src,pts_dst=pts_dst)
         total_max_overlap = self._max_pixel_overlap(self._MASK_COURT_EDGES,pts_src,pts_dst=pts_dst)
-        goodness = float(np.count_nonzero(mapped_edge_img > 100)) / total_max_overlap
+        if total_max_overlap != 0:
+            goodness = float(np.count_nonzero(mapped_edge_img > 100)) / total_max_overlap
+        else:
+            goodness = 0
+
         return goodness
 
     def _get_four_intersections(self,l1:list,l2:list,l3:list,l4:list,relax_factor=0):

src/processrunner.py

@@ -9,7 +9,7 @@ class ProcessRunner:
     Performs player, team, shot, and courtline detection in sequence.
     Effect: updates GameState with statistics and produces courtline video.
     """
-    def __init__(self, video_path, players_tracking, ball_tracking, output_video_path, 
+    def __init__(self, video_path, players_tracking, ball_tracking, output_video_path,
                  output_video_path_reenc) -> None:
         self.video_path = video_path
         self.players_tracking = players_tracking
@@ -67,7 +67,7 @@ def run_video_render(self):
         """Runs video rendering and reencodes, stores to output_video_path_reenc."""
         videoRender = video_render.VideoRender(self.homography)
         videoRender.render_video(self.state.states, self.state.players, self.output_video_path)
-        videoRender.reencode(self.output_video_path, 
+        videoRender.reencode(self.output_video_path,
                               self.output_video_path_reenc)