Pose Estimation

- Create branch
- Write src/pose_estimation/pose_estimate.py
- Works, but does not do anything right now
- When running, terminal outputs non fatal errors (?), need to test

.gitignore

@@ -1,3 +1,16 @@
 venv
 .env
 __pycache__
+# Ignore specific large files
+yolov8m-pose.pt
+*.mp4  # ignores all .mp4 files; remove this line if some .mp4 files should be tracked
+ball/lib/python3.11/site-packages/torch/lib/libtorch_cpu.dylib
+venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu.dylib
+
+# Ignore entire directories
+ball/
+src/pose_estimation/
+tmp/
+
+# Ignore all .json files in tmp/ directory; adjust as per requirement
+tmp/*.json

src/main.py

@@ -36,7 +36,7 @@ def main(video_path):
 
     modelrunner = ModelRunner(video_path, model_vars)
     modelrunner.run()
-    people_output, ball_output = modelrunner.fetch_output()
+    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'
 

src/modelrunner.py

@@ -6,6 +6,7 @@
 import pickle
 import subprocess
 from typing import Tuple
+from pose_estimation.pose_estimate import PoseEstimator
 
 class ModelRunner:
     """
@@ -15,6 +16,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:
@@ -50,16 +52,16 @@ def run(self):
         # subprocess.run(['bash', 'src/StrongSORT-YOLO/run_tracker.sh'])
         with open('tmp/output.pickle', 'rb') as f:
             self.output_dict = pickle.load(f)
+        self.pose_estimator.estimate_pose()
 
-
-    def fetch_output(self) -> Tuple[str, str]:
+    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 +74,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/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