diff --git a/.gitignore b/.gitignore index 5dd8186b..620c344c 100644 --- a/.gitignore +++ b/.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 \ No newline at end of file diff --git a/src/main.py b/src/main.py index 966ac9ad..407b7535 100644 --- a/src/main.py +++ b/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' diff --git a/src/modelrunner.py b/src/modelrunner.py index a9c3444c..858fb90d 100644 --- a/src/modelrunner.py +++ b/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' diff --git a/src/processing/courtline_detect.py b/src/processing/courtline_detect.py index 4bc7aa13..44a12719 100644 --- a/src/processing/courtline_detect.py +++ b/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): diff --git a/src/processing/pose_estimate.py b/src/processing/pose_estimate.py deleted file mode 100644 index c3da9538..00000000 --- a/src/processing/pose_estimate.py +++ /dev/null @@ -1,5 +0,0 @@ -from ultralytics import YOLO - -model = YOLO('yolov8m-pose.pt') - -results = model(source = "test_video.mp4", show = True, conf = 0.3, save = True) \ No newline at end of file diff --git a/src/processrunner.py b/src/processrunner.py index 6c3b4365..7f450921 100644 --- a/src/processrunner.py +++ b/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 diff --git a/tmp/court_video.mp4 b/tmp/court_video.mp4 index 33a4bc47..faacb1c6 100644 Binary files a/tmp/court_video.mp4 and b/tmp/court_video.mp4 differ