create new branch to implement pose estimation

CornellDataScience · Oct 1, 2023 · ddf658c · ddf658c
1 parent 80eb10a
commit ddf658c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/.DS_Store b/.DS_Store
diff --git a/src/main.py b/src/main.py
@@ -40,7 +40,7 @@ def main(video_path):
     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()

diff --git a/src/processing/pose_estimate.py b/src/processing/pose_estimate.py
@@ -0,0 +1,5 @@
+from ultralytics import YOLO
+
+model = YOLO('yolov8m-pose.pt')
+
+results = model(source = "test_video.mp4", show = True, conf = 0.3, save = True)