From ecfa8cafd99d52824cca888e32c6645a57dee7b1 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 14:30:09 -0600 Subject: [PATCH 01/21] move tracker dependencies to extras --- pyproject.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2d6f6dc..2469426 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,8 +57,6 @@ keywords = [ dependencies = [ "opencv-contrib-python==4.8.*", "pydantic==1.*", - "mediapipe==0.10.9", - "ultralytics~=8.0.202", ] requires-python = ">=3.9,<3.12" @@ -66,6 +64,9 @@ dynamic = ["version", "description"] [project.optional-dependencies] dev = ["black", "bumpver", "isort", "pip-tools", "pytest"] +mediapipe = ["mediapipe==0.10.9"] +yolo = ["ultralytics~=8.0.202"] +all = ["ultralytics~=8.0.202", "mediapipe==0.10.9"] [project.urls] Homepage = "https://github.com/freemocap/skellytracker" From b067b8aa70164bc1e383a3daeb22f5b10425a8d2 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 16:01:03 -0600 Subject: [PATCH 02/21] add tqdm --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 2469426..8b8fd04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ keywords = [ dependencies = [ "opencv-contrib-python==4.8.*", "pydantic==1.*", + "tqdm==4.*", ] requires-python = ">=3.9,<3.12" From 7d9e2aa33015bda984f43426936f0009076b9c01 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 16:03:52 -0600 Subject: [PATCH 03/21] support optional imports in init --- skellytracker/__init__.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index e9caecd..3a9c9d2 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -23,9 +23,18 @@ from skellytracker.system.default_paths import get_log_file_path from skellytracker.system.logging_configuration import configure_logging -from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import MediapipeHolisticTracker -from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker -from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import YOLOMediapipeComboTracker +try: + from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import MediapipeHolisticTracker +except: + print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") +try: + from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker +except: + print("To use yolo_tracker, install skellytracker[yolo]") +try: + from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import YOLOMediapipeComboTracker +except: + print("To use yolo_mediapipe_combo_tracker, install skellytracker[mediapipe, yolo] or skellytracker[all]") From 75905a3baa21a808520e84c9d2c17b5b42b7c95e Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 19:16:41 -0600 Subject: [PATCH 04/21] conditional imports in RUN_ME --- skellytracker/RUN_ME.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index 66aa9e1..c5eea5f 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -1,19 +1,26 @@ import cv2 + from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import ( BrightestPointTracker, ) from skellytracker.trackers.charuco_tracker.charuco_tracker import CharucoTracker -from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( - MediapipeHolisticTracker, -) from skellytracker.trackers.segment_anything_tracker.segment_anything_tracker import ( SAMTracker, ) -from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker -from skellytracker.trackers.yolo_object_tracker.yolo_object_tracker import ( - YOLOObjectTracker, -) +try: + from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( + MediapipeHolisticTracker, + ) +except: + print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") +try: + from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker + from skellytracker.trackers.yolo_object_tracker.yolo_object_tracker import ( + YOLOObjectTracker, + ) +except: + print("To use yolo_tracker, install skellytracker[yolo]") def main(demo_tracker: str = "mediapipe_holistic_tracker"): From cc8d56cd616179bb580dd793e5c1fab3eeed07d6 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 19:21:22 -0600 Subject: [PATCH 05/21] more conditional imports --- skellytracker/RUN_ME.py | 6 ++--- skellytracker/process_folder_of_videos.py | 31 +++++++++++++++-------- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index c5eea5f..7306071 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -5,9 +5,6 @@ BrightestPointTracker, ) from skellytracker.trackers.charuco_tracker.charuco_tracker import CharucoTracker -from skellytracker.trackers.segment_anything_tracker.segment_anything_tracker import ( - SAMTracker, -) try: from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( MediapipeHolisticTracker, @@ -19,6 +16,9 @@ from skellytracker.trackers.yolo_object_tracker.yolo_object_tracker import ( YOLOObjectTracker, ) + from skellytracker.trackers.segment_anything_tracker.segment_anything_tracker import ( + SAMTracker, + ) except: print("To use yolo_tracker, install skellytracker[yolo]") diff --git a/skellytracker/process_folder_of_videos.py b/skellytracker/process_folder_of_videos.py index c32a2b6..aaa60e7 100644 --- a/skellytracker/process_folder_of_videos.py +++ b/skellytracker/process_folder_of_videos.py @@ -2,7 +2,7 @@ from multiprocessing import Pool, cpu_count from pathlib import Path import sys -from typing import Any, Optional, Type +from typing import Optional import numpy as np from pydantic import BaseModel @@ -11,17 +11,26 @@ from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import ( BrightestPointTracker, ) -from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( - MediapipeHolisticTracker, -) -from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import ( - YOLOMediapipeComboTracker, -) -from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker -from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import ( - MediapipeTrackingParams, -) from skellytracker.utilities.get_video_paths import get_video_paths +try: + from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import ( + YOLOMediapipeComboTracker, + ) +except: + print("To use yolo_mediapipe_combo_tracker, install skellytracker[yolo, mediapipe]") +try: + from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker +except: + print("To use yolo_tracker, install skellytracker[yolo]") +try: + from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( + MediapipeHolisticTracker, + ) + from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import ( + MediapipeTrackingParams, + ) +except: + print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") logger = logging.getLogger(__name__) From 4bb84085e5a621d907b2db82a3ea907bd15590fb Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 19:53:05 -0600 Subject: [PATCH 06/21] parameterize params in process_folder_of_videos --- skellytracker/process_folder_of_videos.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/skellytracker/process_folder_of_videos.py b/skellytracker/process_folder_of_videos.py index aaa60e7..fccb05b 100644 --- a/skellytracker/process_folder_of_videos.py +++ b/skellytracker/process_folder_of_videos.py @@ -20,6 +20,7 @@ print("To use yolo_mediapipe_combo_tracker, install skellytracker[yolo, mediapipe]") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker + from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOTrackingParams except: print("To use yolo_tracker, install skellytracker[yolo]") try: @@ -48,7 +49,7 @@ def process_folder_of_videos( synchronized_video_path: Path, output_folder_path: Optional[Path] = None, annotated_video_path: Optional[Path] = None, - num_processes: int = None, + num_processes: Optional[int] = None, ) -> np.ndarray: """ Process a folder of synchronized videos with the given tracker. @@ -176,6 +177,18 @@ def get_tracker(tracker_name: str, tracking_params: BaseModel) -> BaseTracker: return tracker +def get_tracker_params(tracker_name: str) -> BaseModel: + if tracker_name == "MediapipeHolisticTracker": + return MediapipeTrackingParams() + elif tracker_name == "YOLOMediapipeComboTracker": + return YOLOTrackingParams() + elif tracker_name == "YOLOPoseTracker": + return YOLOTrackingParams() + elif tracker_name == "BrightestPointTracker": + return BaseModel() + else: + raise ValueError("Invalid tracker type") + if __name__ == "__main__": synchronized_video_path = Path( @@ -186,7 +199,7 @@ def get_tracker(tracker_name: str, tracking_params: BaseModel) -> BaseTracker: process_folder_of_videos( tracker_name=tracker_name, - tracking_params=MediapipeTrackingParams(), + tracking_params=get_tracker_params(tracker_name=tracker_name), synchronized_video_path=synchronized_video_path, num_processes=num_processes, ) From b4ded949ebfb68916a6173668af0617472b60f5b Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 18 Mar 2024 21:36:10 -0600 Subject: [PATCH 07/21] conditional imports in single image run --- skellytracker/SINGLE_IMAGE_RUN.py | 40 +++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index a301b11..ec9bf4c 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -1,10 +1,22 @@ import cv2 from pathlib import Path -from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import BrightestPointTracker +from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import ( + BrightestPointTracker, +) from skellytracker.trackers.charuco_tracker.charuco_tracker import CharucoTracker -from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import MediapipeHolisticTracker -from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker + +try: + from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( + MediapipeHolisticTracker, + ) +except: + print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") +try: + from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker +except: + print("To use yolo_tracker, install skellytracker[yolo]") + if __name__ == "__main__": demo_tracker = "brightest_point_tracker" @@ -14,16 +26,20 @@ BrightestPointTracker().image_demo(image_path=image_path) elif demo_tracker == "charuco_tracker": - CharucoTracker(squaresX=7, - squaresY=5, - dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250)).image_demo(image_path=image_path) + CharucoTracker( + squaresX=7, + squaresY=5, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ).image_demo(image_path=image_path) elif demo_tracker == "mediapipe_holistic_tracker": - MediapipeHolisticTracker(model_complexity=2, - min_detection_confidence=0.5, - min_tracking_confidence=0.5, - static_image_mode=False, - smooth_landmarks=True).image_demo(image_path=image_path) + MediapipeHolisticTracker( + model_complexity=2, + min_detection_confidence=0.5, + min_tracking_confidence=0.5, + static_image_mode=False, + smooth_landmarks=True, + ).image_demo(image_path=image_path) elif demo_tracker == "yolo_tracker": - YOLOPoseTracker(model_size="high_res").image_demo(image_path=image_path) \ No newline at end of file + YOLOPoseTracker(model_size="high_res").image_demo(image_path=image_path) From 775be98eb14b5379c70e2356d888d91ec2dbe9c2 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Wed, 24 Apr 2024 23:08:03 -0600 Subject: [PATCH 08/21] Update skellytracker/SINGLE_IMAGE_RUN.py Co-authored-by: jonmatthis --- skellytracker/SINGLE_IMAGE_RUN.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index ec9bf4c..f4e8a5a 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -11,7 +11,7 @@ MediapipeHolisticTracker, ) except: - print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") + print("\n\nTo use mediapipe_holistic_tracker, install skellytracker[mediapipe]\n\n") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker except: From 5e98b07d6fcf26faaea22643aa3cc07702d214f9 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Wed, 24 Apr 2024 23:08:40 -0600 Subject: [PATCH 09/21] Apply suggestions from code review Co-authored-by: jonmatthis --- skellytracker/SINGLE_IMAGE_RUN.py | 2 +- skellytracker/process_folder_of_videos.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index f4e8a5a..9b3dea9 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -15,7 +15,7 @@ try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker except: - print("To use yolo_tracker, install skellytracker[yolo]") + print("\n\nTo use yolo_tracker, install skellytracker[yolo]\n\n") if __name__ == "__main__": diff --git a/skellytracker/process_folder_of_videos.py b/skellytracker/process_folder_of_videos.py index fccb05b..12f1bc9 100644 --- a/skellytracker/process_folder_of_videos.py +++ b/skellytracker/process_folder_of_videos.py @@ -17,7 +17,7 @@ YOLOMediapipeComboTracker, ) except: - print("To use yolo_mediapipe_combo_tracker, install skellytracker[yolo, mediapipe]") + print("\n\nTo use yolo_mediapipe_combo_tracker, install skellytracker[yolo, mediapipe]\n\n") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOTrackingParams From 97cfa098a59d12192f81c1a8e99dcf529a35808e Mon Sep 17 00:00:00 2001 From: philipqueen Date: Thu, 25 Apr 2024 11:24:15 -0600 Subject: [PATCH 10/21] Bump version v2024.03.1013 -> v2024.04.1014 --- pyproject.toml | 2 +- skellytracker/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 10d0e57..97ab4c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ all = ["ultralytics~=8.0.202", "mediapipe==0.10.9"] Homepage = "https://github.com/freemocap/skellytracker" [tool.bumpver] -current_version = "v2024.03.1013" +current_version = "v2024.04.1014" version_pattern = "vYYYY.0M.BUILD[-TAG]" commit_message = "Bump version {old_version} -> {new_version}" diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index e8543f1..54507c5 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -1,7 +1,7 @@ """Top-level package for skellytracker""" __package_name__ = "skellytracker" -__version__ = "v2024.03.1013" +__version__ = "v2024.04.1014" __author__ = """Skelly FreeMoCap""" __email__ = "info@freemocap.org" From 79d7a117fc9759824dddc69e6cb70dba29cc767f Mon Sep 17 00:00:00 2001 From: philipqueen Date: Fri, 7 Jun 2024 11:41:13 -0600 Subject: [PATCH 11/21] Patch yolo object cpu conversion (#37) * patch CPU conversion error on YOLO object tracker * specify to run on cpu * specify cpu in yolo mediapipe combo tracker * remove device flag for yolo object tracker --- .../yolo_mediapipe_combo_tracker.py | 2 +- .../trackers/yolo_object_tracker/yolo_object_tracker.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py b/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py index d6ff664..e0cb3cb 100644 --- a/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py +++ b/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py @@ -55,7 +55,7 @@ def __init__( def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject]: yolo_results = self.model(image, classes=0, max_det=1, verbose=False) - box_xyxy = np.asarray(yolo_results[0].boxes.xyxy).flatten() + box_xyxy = np.asarray(yolo_results[0].boxes.xyxy.cpu()).flatten() if box_xyxy.size > 0: box_left, box_top, box_right, box_bottom = box_xyxy diff --git a/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py b/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py index 0cfa29a..62dca7b 100644 --- a/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py +++ b/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py @@ -32,7 +32,7 @@ def __init__( def process_image(self, image, **kwargs) -> Dict[str, TrackedObject]: results = self.model(image, classes=self.classes, max_det=1, verbose=False, conf=self.confidence_threshold) - box_xyxy = np.asarray(results[0].boxes.xyxy).flatten() + box_xyxy = np.asarray(results[0].boxes.xyxy.cpu()).flatten() # if on GPU, need to copy to CPU before np array conversion if box_xyxy.size > 0: self.tracked_objects["object"].pixel_x = (box_xyxy[0] + box_xyxy[2]) / 0.5 From 79959ad1a91653e6b24579263db8259c45912381 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Fri, 7 Jun 2024 11:49:41 -0600 Subject: [PATCH 12/21] Bump version v2024.04.1014 -> v2024.06.1015 --- pyproject.toml | 2 +- skellytracker/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 97ab4c4..68b0ab4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ all = ["ultralytics~=8.0.202", "mediapipe==0.10.9"] Homepage = "https://github.com/freemocap/skellytracker" [tool.bumpver] -current_version = "v2024.04.1014" +current_version = "v2024.06.1015" version_pattern = "vYYYY.0M.BUILD[-TAG]" commit_message = "Bump version {old_version} -> {new_version}" diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index 54507c5..3b7c967 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -1,7 +1,7 @@ """Top-level package for skellytracker""" __package_name__ = "skellytracker" -__version__ = "v2024.04.1014" +__version__ = "v2024.06.1015" __author__ = """Skelly FreeMoCap""" __email__ = "info@freemocap.org" From ef0c0ff99607e09f9458432e2840387cbc1613e3 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 17 Jun 2024 12:19:30 -0600 Subject: [PATCH 13/21] Test pose estimation trackers (#39) --- .github/workflows/python-testing.yml | 41 ++++++ pyproject.toml | 1 + skellytracker/__init__.py | 4 +- skellytracker/system/default_paths.py | 1 + skellytracker/test/__init__.py | 0 .../tests/__init__.py | 0 skellytracker/tests/conftest.py | 14 ++ .../tests/test_mediapipe_holistic_tracker.py | 126 ++++++++++++++++++ skellytracker/{test => tests}/test_test.py | 0 skellytracker/tests/test_yolo_pose_tracker.py | 117 ++++++++++++++++ .../trackers/base_tracker/base_recorder.py | 4 +- .../trackers/base_tracker/base_tracker.py | 1 - .../brightest_point_tracker.py | 2 - .../charuco_tracker/charuco_tracker.py | 2 - .../mediapipe_holistic_recorder.py | 2 +- .../mediapipe_holistic_tracker.py | 13 +- .../trackers/yolo_tracker/yolo_tracker.py | 14 +- .../utilities/download_test_image.py | 32 +++++ 18 files changed, 349 insertions(+), 25 deletions(-) create mode 100644 .github/workflows/python-testing.yml delete mode 100644 skellytracker/test/__init__.py rename __init__.py => skellytracker/tests/__init__.py (100%) create mode 100644 skellytracker/tests/conftest.py create mode 100644 skellytracker/tests/test_mediapipe_holistic_tracker.py rename skellytracker/{test => tests}/test_test.py (100%) create mode 100644 skellytracker/tests/test_yolo_pose_tracker.py create mode 100644 skellytracker/utilities/download_test_image.py diff --git a/.github/workflows/python-testing.yml b/.github/workflows/python-testing.yml new file mode 100644 index 0000000..9f5315e --- /dev/null +++ b/.github/workflows/python-testing.yml @@ -0,0 +1,41 @@ +name: SkellyTracker Tests + +on: + pull_request: + branches: [ main ] + paths: + - 'skellytracker/**' + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: System Info + run: | + uname -a || true + lsb_release -a || true + gcc --version || true + env + - name: Set up Python 3.x + uses: actions/setup-python@v4 + with: + # Semantic version range syntax or exact version of a Python version + python-version: '3.9' + # Optional - x64 or x86 architecture, defaults to x64 + architecture: 'x64' + cache: 'pip' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install "-e.[all]" + - name: Fix OpenCV conflict + run: | + pip uninstall -y opencv-python opencv-contrib-python + pip install opencv-contrib-python==4.8.1.78 + - name: Set PYTHONPATH + run: echo "PYTHONPATH=$PYTHONPATH:$(pwd)/skellytracker" >> $GITHUB_ENV + - name: Run Tests with Pytest + run: | + pip install pytest + pytest skellytracker/tests \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 68b0ab4..e4bd6df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ keywords = [ #dynamic = ["dependencies"] dependencies = [ "opencv-contrib-python==4.8.*", + "numpy<2", "pydantic==1.*", "tqdm==4.*", ] diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index 3b7c967..3b8046c 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -20,6 +20,8 @@ print(f"adding base_package_path: {base_package_path} : to sys.path") sys.path.insert(0, str(base_package_path)) # add parent directory to sys.path +print(f"sys path: {sys.path}") + from skellytracker.system.default_paths import get_log_file_path from skellytracker.system.logging_configuration import configure_logging @@ -38,4 +40,4 @@ -configure_logging(log_file_path=get_log_file_path()) \ No newline at end of file +configure_logging(log_file_path=str(get_log_file_path())) \ No newline at end of file diff --git a/skellytracker/system/default_paths.py b/skellytracker/system/default_paths.py index 37bd5f6..58a5d4c 100644 --- a/skellytracker/system/default_paths.py +++ b/skellytracker/system/default_paths.py @@ -7,6 +7,7 @@ BASE_FOLDER_NAME = f"{__package_name__}_data" LOGS_INFO_AND_SETTINGS_FOLDER_NAME = "logs_info_and_settings" LOG_FILE_FOLDER_NAME = "logs" +FIGSHARE_TEST_IMAGE_URL = "https://figshare.com/ndownloader/files/47043898" def get_base_folder_path(): base_folder = Path().home() / BASE_FOLDER_NAME diff --git a/skellytracker/test/__init__.py b/skellytracker/test/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/__init__.py b/skellytracker/tests/__init__.py similarity index 100% rename from __init__.py rename to skellytracker/tests/__init__.py diff --git a/skellytracker/tests/conftest.py b/skellytracker/tests/conftest.py new file mode 100644 index 0000000..3f226dd --- /dev/null +++ b/skellytracker/tests/conftest.py @@ -0,0 +1,14 @@ +import numpy as np +import pytest +from skellytracker.utilities.download_test_image import download_test_image + + +class SessionInfo: + test_image: np.ndarray + +def pytest_sessionstart(session): + SessionInfo.test_image = download_test_image() + +@pytest.fixture +def test_image(): + return SessionInfo.test_image \ No newline at end of file diff --git a/skellytracker/tests/test_mediapipe_holistic_tracker.py b/skellytracker/tests/test_mediapipe_holistic_tracker.py new file mode 100644 index 0000000..58026a5 --- /dev/null +++ b/skellytracker/tests/test_mediapipe_holistic_tracker.py @@ -0,0 +1,126 @@ +import math +import cv2 +import pytest +import numpy as np + + +from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import ( + MediapipeModelInfo, +) +from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( + MediapipeHolisticTracker, +) + + +@pytest.mark.usefixtures("test_image") +def test_process_image(test_image): + tracker = MediapipeHolisticTracker(model_complexity=0) + tracked_objects = tracker.process_image(test_image) + + assert len(tracked_objects) == 4 + assert tracked_objects["pose_landmarks"] is not None + assert tracked_objects["pose_landmarks"].extra["landmarks"] is not None + assert tracked_objects["right_hand_landmarks"] is not None + assert tracked_objects["right_hand_landmarks"].extra["landmarks"] is not None + assert tracked_objects["left_hand_landmarks"] is not None + assert tracked_objects["left_hand_landmarks"].extra["landmarks"] is not None + assert tracked_objects["face_landmarks"] is not None + assert tracked_objects["face_landmarks"].extra["landmarks"] is not None + + +@pytest.mark.usefixtures("test_image") +def test_annotate_image(test_image): + tracker = MediapipeHolisticTracker(model_complexity=0) + tracker.process_image(test_image) + + assert tracker.annotated_image is not None + + +@pytest.mark.usefixtures("test_image") +def test_record(test_image): + tracker = MediapipeHolisticTracker(model_complexity=0) + tracked_objects = tracker.process_image(test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + assert len(tracker.recorder.recorded_objects[0]) == 4 + + processed_results = tracker.recorder.process_tracked_objects( + image_size=test_image.shape[:2] + ) + assert processed_results is not None + assert processed_results.shape == ( + 1, + MediapipeModelInfo.num_tracked_points_total, + 3, + ) + + expected_results = np.array( + [ + [ + [724.9490356445312, 80.74257552623749, -993.9854431152344], + [748.7973022460938, 76.24467730522156, -955.4129028320312], + [757.65625, 77.80313193798065, -955.689697265625], + [765.3969573974609, 79.39812362194061, -955.8861541748047], + [724.6346282958984, 73.26869666576385, -956.6817474365234], + [717.9697418212891, 72.75395393371582, -956.7975616455078], + [711.0582733154297, 72.3196667432785, -956.7622375488281], + [775.3928375244141, 86.46261692047119, -668.9157867431641], + [700.8257293701172, 76.53615295886993, -673.5513305664062], + [736.4148712158203, 91.9196355342865, -881.9746398925781], + [704.9551391601562, 87.58252501487732, -883.7062072753906], + [822.0133209228516, 152.90948510169983, -438.8511276245117], + [543.9170837402344, 128.66411805152893, -474.09820556640625], + [830.2363586425781, 222.54773139953613, -348.53546142578125], + [391.7664337158203, 185.8674716949463, -379.22679901123047], + [850.4001617431641, 281.42955780029297, -597.9225921630859], + [251.3930320739746, 234.69798803329468, -603.9112854003906], + [856.3973236083984, 297.1209526062012, -668.1196594238281], + [213.4073257446289, 246.7552900314331, -670.8364105224609], + [846.2405395507812, 297.8205370903015, -768.86962890625], + [218.36742401123047, 250.46862602233887, -770.2429962158203], + [841.2351226806641, 292.3808026313782, -645.0504302978516], + [237.8342628479004, 244.97743606567383, -649.0667724609375], + [642.2610473632812, 267.2499203681946, 19.359580278396606], + [489.5075225830078, 252.38561153411865, -18.99002432823181], + [545.5322647094727, 353.45691204071045, 95.11228561401367], + [421.98192596435547, 343.2889795303345, -111.5739917755127], + [483.64437103271484, 413.63812923431396, 892.3109436035156], + [386.21551513671875, 410.32382011413574, 586.4739990234375], + [469.1791534423828, 416.32561683654785, 958.9944458007812], + [386.70398712158203, 417.09611892700195, 642.3279571533203], + [473.2674789428711, 447.85693645477295, 729.4257354736328], + [355.9325408935547, 444.5236587524414, 379.90428924560547], + [253.68976593017578, 234.29851055145264, 0.0001452891228836961], + [240.58324813842773, 234.31915283203125, -17.537919282913208], + [225.88279724121094, 236.56160831451416, -26.281862258911133], + [216.74800872802734, 239.36949491500854, -33.34794282913208], + [207.5124740600586, 241.5495729446411, -39.871764183044434], + [203.52766036987305, 250.0307822227478, -13.84335994720459], + [184.57786560058594, 257.0957851409912, -22.172749042510986], + [172.95087814331055, 261.3565707206726, -29.567382335662842], + [163.63656997680664, 264.60676431655884, -34.58906173706055], + [210.67567825317383, 252.71584510803223, -10.44108510017395], + [192.9056167602539, 260.5885577201843, -17.135307788848877], + [181.97023391723633, 265.8717155456543, -23.215365409851074], + [173.35857391357422, 269.7653818130493, -27.345290184020996], + [219.38091278076172, 253.6610770225525, -8.99298369884491], + [204.17118072509766, 261.2510418891907, -14.699053764343262], + [195.27652740478516, 266.1923360824585, -19.270341396331787], + [187.59790420532227, 269.8603105545044, -22.399024963378906], + [230.7835578918457, 253.28086853027344, -9.420499801635742], + [223.0868148803711, 259.03817653656006, -14.548875093460083], + [218.3838653564453, 262.8340816497803, -17.04281210899353], + [214.09095764160156, 265.90606927871704, -18.470606803894043], + [831.3526916503906, 290.13811111450195, 0.00027009031327906996], + [834.6914672851562, 290.1472735404968, -28.73823642730713], + [838.1427764892578, 291.96690559387207, -45.007004737854004], + [837.1710205078125, 295.0093674659729, -56.84781551361084], + [838.1177520751953, 298.5115385055542, -67.84458637237549], + [836.1019134521484, 307.75739192962646, -34.55303907394409], + [831.5267944335938, 317.115318775177, -48.21352005004883], + ] + ] + ) + assert np.allclose( + processed_results[:, :60, :], expected_results[:, :60, :], atol=1 + ) diff --git a/skellytracker/test/test_test.py b/skellytracker/tests/test_test.py similarity index 100% rename from skellytracker/test/test_test.py rename to skellytracker/tests/test_test.py diff --git a/skellytracker/tests/test_yolo_pose_tracker.py b/skellytracker/tests/test_yolo_pose_tracker.py new file mode 100644 index 0000000..b0c5343 --- /dev/null +++ b/skellytracker/tests/test_yolo_pose_tracker.py @@ -0,0 +1,117 @@ +import math +import pytest +import numpy as np + + +from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOModelInfo +from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker + + +@pytest.mark.usefixtures("test_image") +def test_process_image(test_image): + tracker = YOLOPoseTracker(model_size="nano") + tracked_objects = tracker.process_image(test_image) + + assert len(tracked_objects) == 1 + tracked_person = tracked_objects["tracked_person"] + assert tracked_person.pixel_x is not None + assert tracked_person.pixel_y is not None + assert math.isclose(tracked_person.pixel_x, 266.48523, rel_tol=1e-2) + assert math.isclose(tracked_person.pixel_y, 273.92798, rel_tol=1e-2) + + landmarks = tracked_person.extra["landmarks"] + assert landmarks is not None + assert landmarks.shape == (1, YOLOModelInfo.num_tracked_points, 2) + + expected_results = np.array( + [ + [ + [392.56927490234375, 140.40118408203125], + [414.94940185546875, 132.90655517578125], + [386.4353332519531, 125.51483154296875], + [446.2061767578125, 157.98883056640625], + [373.6619873046875, 138.93646240234375], + [453.78662109375, 265.081787109375], + [317.9375305175781, 231.9653778076172], + [465.893310546875, 396.12274169921875], + [220.12176513671875, 325.96636962890625], + [465.23358154296875, 499.0487365722656], + [142.17066955566406, 407.7397155761719], + [352.325439453125, 468.44671630859375], + [268.8867492675781, 448.7227783203125], + [310.899658203125, 630.5478515625], + [227.7810821533203, 617.9011840820312], + [269.08587646484375, 733.4285888671875], + [213.1557159423828, 741.54541015625], + ] + ] + ) + + assert np.allclose(landmarks, expected_results) + + +@pytest.mark.usefixtures("test_image") +def test_annotate_image(test_image): + tracker = YOLOPoseTracker(model_size="nano") + tracker.process_image(test_image) + + assert tracker.annotated_image is not None + + +@pytest.mark.usefixtures("test_image") +def test_record(test_image): + tracker = YOLOPoseTracker(model_size="nano") + tracked_objects = tracker.process_image(test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + + processed_results = tracker.recorder.process_tracked_objects() + assert processed_results is not None + assert processed_results.shape == (1, YOLOModelInfo.num_tracked_points, 3) + + expected_results = np.array( + [ + [ + [392.56927490234375, 140.40118408203125, np.nan], + [414.94940185546875, 132.90655517578125, np.nan], + [386.4353332519531, 125.51483154296875, np.nan], + [446.2061767578125, 157.98883056640625, np.nan], + [373.6619873046875, 138.93646240234375, np.nan], + [453.78662109375, 265.081787109375, np.nan], + [317.9375305175781, 231.9653778076172, np.nan], + [465.893310546875, 396.12274169921875, np.nan], + [220.12176513671875, 325.96636962890625, np.nan], + [465.23358154296875, 499.0487365722656, np.nan], + [142.17066955566406, 407.7397155761719, np.nan], + [352.325439453125, 468.44671630859375, np.nan], + [268.8867492675781, 448.7227783203125, np.nan], + [310.899658203125, 630.5478515625, np.nan], + [227.7810821533203, 617.9011840820312, np.nan], + [269.08587646484375, 733.4285888671875, np.nan], + [213.1557159423828, 741.54541015625, np.nan], + ] + ] + ) + assert np.allclose(processed_results[:, :, :2], expected_results[:, :, :2], atol=1e-2) + assert np.isnan(processed_results[:, :, 2]).all() + + +class MockKeypoints: + xy: list = [] + + +class MockResults: + keypoints: MockKeypoints = MockKeypoints() + + +def test_unpack_empty_results(): + tracker = YOLOPoseTracker(model_size="nano") + results = [MockResults()] + tracker.unpack_results(results=results) + + tracked_person = tracker.tracked_objects["tracked_person"] + assert tracked_person.pixel_x is None + assert tracked_person.pixel_y is None + assert tracked_person.extra["landmarks"].shape == (1, YOLOModelInfo.num_tracked_points, 2) + assert np.isnan(tracked_person.extra["landmarks"][0, :, 0]).all() + assert np.isnan(tracked_person.extra["landmarks"][0, :, 1]).all() diff --git a/skellytracker/trackers/base_tracker/base_recorder.py b/skellytracker/trackers/base_tracker/base_recorder.py index f81f0b3..9974da0 100644 --- a/skellytracker/trackers/base_tracker/base_recorder.py +++ b/skellytracker/trackers/base_tracker/base_recorder.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod import logging -from typing import Dict +from typing import Dict, Optional import numpy as np @@ -20,7 +20,7 @@ def __init__(self): @abstractmethod def record( - self, tracked_objects: Dict[str, TrackedObject], annotated_image: np.ndarray + self, tracked_objects: Dict[str, TrackedObject], annotated_image: Optional[np.ndarray] = None ) -> None: """ Record the tracked objects as they are created by the tracker. diff --git a/skellytracker/trackers/base_tracker/base_tracker.py b/skellytracker/trackers/base_tracker/base_tracker.py index 946b576..a3ab1c5 100644 --- a/skellytracker/trackers/base_tracker/base_tracker.py +++ b/skellytracker/trackers/base_tracker/base_tracker.py @@ -31,7 +31,6 @@ def __init__( ): self.recorder = recorder self.annotated_image = None - self.raw_image = None self.tracked_objects: Dict[str, TrackedObject] = {} for name in tracked_object_names: diff --git a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py index fb573ec..d934963 100644 --- a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py +++ b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py @@ -53,8 +53,6 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] self.tracked_objects["brightest_point"].pixel_y = largest_patch_centroid[1] self.tracked_objects["brightest_point"].extra["thresholded_image"] = thresholded_image - self.raw_image = image.copy() - self.annotated_image = self.annotate_image(image=image, tracked_objects=self.tracked_objects) diff --git a/skellytracker/trackers/charuco_tracker/charuco_tracker.py b/skellytracker/trackers/charuco_tracker/charuco_tracker.py index 408a988..c0d172c 100644 --- a/skellytracker/trackers/charuco_tracker/charuco_tracker.py +++ b/skellytracker/trackers/charuco_tracker/charuco_tracker.py @@ -44,8 +44,6 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] self.tracked_objects[object_id].pixel_x = corner[0][0] self.tracked_objects[object_id].pixel_y = corner[0][1] - self.raw_image = image.copy() - self.annotated_image = self.annotate_image(image=image, tracked_objects=self.tracked_objects) diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py index 4e70b67..9e0f298 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py @@ -21,7 +21,7 @@ def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: def process_tracked_objects(self, **kwargs) -> np.ndarray: image_size = kwargs.get("image_size") if image_size is None: - raise ValueError("image_size must be provided to process tracked objects") + raise ValueError(f"image_size must be provided to process tracked objects from {__class__.__name__}") self.recorded_objects_array = np.zeros( ( len(self.recorded_objects), diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py index d5a97f3..84c34c6 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py @@ -57,8 +57,6 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] "landmarks" ] = results.right_hand_landmarks - self.raw_image = image.copy() - self.annotated_image = self.annotate_image( image=image, tracked_objects=self.tracked_objects ) @@ -68,29 +66,30 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] def annotate_image( self, image: np.ndarray, tracked_objects: Dict[str, TrackedObject], **kwargs ) -> np.ndarray: + annotated_image = image.copy() # Draw the pose, face, and hand landmarks on the image self.mp_drawing.draw_landmarks( - image, + annotated_image, tracked_objects["pose_landmarks"].extra["landmarks"], self.mp_holistic.POSE_CONNECTIONS, ) self.mp_drawing.draw_landmarks( - image, + annotated_image, tracked_objects["face_landmarks"].extra["landmarks"], self.mp_holistic.FACEMESH_TESSELATION, ) self.mp_drawing.draw_landmarks( - image, + annotated_image, tracked_objects["left_hand_landmarks"].extra["landmarks"], self.mp_holistic.HAND_CONNECTIONS, ) self.mp_drawing.draw_landmarks( - image, + annotated_image, tracked_objects["right_hand_landmarks"].extra["landmarks"], self.mp_holistic.HAND_CONNECTIONS, ) - return image + return annotated_image if __name__ == "__main__": diff --git a/skellytracker/trackers/yolo_tracker/yolo_tracker.py b/skellytracker/trackers/yolo_tracker/yolo_tracker.py index 760c64c..8bce06b 100644 --- a/skellytracker/trackers/yolo_tracker/yolo_tracker.py +++ b/skellytracker/trackers/yolo_tracker/yolo_tracker.py @@ -21,26 +21,22 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] self.unpack_results(results) - self.annotated_image = self.annotate_image(image, results=results, **kwargs) + self.annotated_image = self.annotate_image(image=image, results=results, **kwargs) return self.tracked_objects - def annotate_image(self, image: np.ndarray, results, **kwargs) -> np.ndarray: + def annotate_image(self, image: np.ndarray, results: list, **kwargs) -> np.ndarray: return results[-1].plot() - def unpack_results(self, results): + def unpack_results(self, results: list): tracked_person = np.asarray(results[-1].keypoints.xy) self.tracked_objects["tracked_person"] = TrackedObject( object_id="tracked_person" ) if tracked_person.size != 0: # add averages of all tracked points as pixel x and y - self.tracked_objects["tracked_person"].pixel_x = np.mean( - tracked_person[:, 0], axis=0 - ) - self.tracked_objects["tracked_person"].pixel_y = np.mean( - tracked_person[:, 1], axis=0 - ) + self.tracked_objects["tracked_person"].pixel_x = np.mean(tracked_person[:, 0]) + self.tracked_objects["tracked_person"].pixel_y = np.mean(tracked_person[:, 1]) self.tracked_objects["tracked_person"].extra["landmarks"] = tracked_person else: self.tracked_objects["tracked_person"].pixel_x = None diff --git a/skellytracker/utilities/download_test_image.py b/skellytracker/utilities/download_test_image.py new file mode 100644 index 0000000..4680827 --- /dev/null +++ b/skellytracker/utilities/download_test_image.py @@ -0,0 +1,32 @@ +import logging +from pathlib import Path + +import cv2 +import numpy as np +import requests + +from skellytracker.system.default_paths import FIGSHARE_TEST_IMAGE_URL + + +logger = logging.getLogger(__name__) + +def download_test_image(test_image_url: str = FIGSHARE_TEST_IMAGE_URL) -> np.ndarray: + try: + logger.info(f"Downloading test image from {test_image_url}...") + + r = requests.get(test_image_url, stream=True, timeout=(5, 60)) + r.raise_for_status() # Check if request was successful + + image_array = np.frombuffer(r.content, np.uint8) + image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) + + logger.info(f"Test image downloaded successfully.") + return image + + except requests.exceptions.RequestException as e: + logger.error(f"Request failed: {e}") + raise e + + +if __name__ == "__main__": + test_data_path = download_test_image() \ No newline at end of file From e70939a248d4a78ff7e91e473e54181cbe85a05a Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 17 Jun 2024 12:28:16 -0600 Subject: [PATCH 14/21] Track multiple brightest points (#38) --- skellytracker/RUN_ME.py | 2 +- skellytracker/SINGLE_IMAGE_RUN.py | 2 +- .../test/test_brightest_point_tracker.py | 84 +++++++++++++++++ .../brightest_point_recorder.py | 19 +++- .../brightest_point_tracker.py | 94 +++++++++++++------ 5 files changed, 164 insertions(+), 37 deletions(-) create mode 100644 skellytracker/test/test_brightest_point_tracker.py diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index 7306071..5a04c8b 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -26,7 +26,7 @@ def main(demo_tracker: str = "mediapipe_holistic_tracker"): if demo_tracker == "brightest_point_tracker": - BrightestPointTracker().demo() + BrightestPointTracker(num_points=2).demo() elif demo_tracker == "charuco_tracker": CharucoTracker( diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index 9b3dea9..f1d337d 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -23,7 +23,7 @@ image_path = Path("/Path/To/Your/Image.jpg") if demo_tracker == "brightest_point_tracker": - BrightestPointTracker().image_demo(image_path=image_path) + BrightestPointTracker(num_points=2).image_demo(image_path=image_path) elif demo_tracker == "charuco_tracker": CharucoTracker( diff --git a/skellytracker/test/test_brightest_point_tracker.py b/skellytracker/test/test_brightest_point_tracker.py new file mode 100644 index 0000000..abe97e6 --- /dev/null +++ b/skellytracker/test/test_brightest_point_tracker.py @@ -0,0 +1,84 @@ +import cv2 +import pytest +import numpy as np + + +from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import ( + BrightestPointTracker, +) + + +@pytest.fixture +def sample_image(): + """ + Create a sample image with bright spots for testing. + """ + image = np.zeros((100, 100, 3), dtype=np.uint8) + cv2.circle(image, (30, 30), 10, (255, 255, 255), -1) # Bright spot + cv2.circle(image, (70, 70), 15, (255, 255, 255), -1) # Brighter spot + return image + + +def test_process_image_with_one_brightest_point(sample_image): + tracker = BrightestPointTracker(num_points=1, luminance_threshold=200) + tracked_objects = tracker.process_image(sample_image) + + assert len(tracked_objects) == 1 + assert tracked_objects["brightest_point_0"].pixel_x == 70 + assert tracked_objects["brightest_point_0"].pixel_y == 70 + + +def test_process_image_with_two_brightest_points(sample_image): + tracker = BrightestPointTracker(num_points=2, luminance_threshold=200) + tracked_objects = tracker.process_image(sample_image) + + assert len(tracked_objects) == 2 + assert tracked_objects["brightest_point_0"].pixel_x == 70 + assert tracked_objects["brightest_point_0"].pixel_y == 70 + assert tracked_objects["brightest_point_1"].pixel_x == 30 + assert tracked_objects["brightest_point_1"].pixel_y == 30 + + +def test_process_image_with_five_brightest_points(sample_image): + tracker = BrightestPointTracker(num_points=5, luminance_threshold=200) + tracked_objects = tracker.process_image(sample_image) + + assert len(tracked_objects) == 5 + assert tracked_objects["brightest_point_0"].pixel_x == 70 + assert tracked_objects["brightest_point_0"].pixel_y == 70 + assert tracked_objects["brightest_point_1"].pixel_x == 30 + assert tracked_objects["brightest_point_1"].pixel_y == 30 + assert tracked_objects["brightest_point_2"].pixel_x is None + assert tracked_objects["brightest_point_2"].pixel_y is None + assert tracked_objects["brightest_point_3"].pixel_x is None + assert tracked_objects["brightest_point_3"].pixel_y is None + assert tracked_objects["brightest_point_4"].pixel_x is None + assert tracked_objects["brightest_point_4"].pixel_y is None + + +def test_annotate_image(sample_image): + tracker = BrightestPointTracker(num_points=2, luminance_threshold=200) + tracked_objects = tracker.process_image(sample_image) + + assert tracker.annotated_image is not None + + # Check if the bright spots have markers + bright_point_0 = tracked_objects["brightest_point_0"] + bright_point_1 = tracked_objects["brightest_point_1"] + + assert tracker.annotated_image[ + bright_point_0.pixel_y, bright_point_0.pixel_x + ].tolist() == [0, 0, 255] + assert tracker.annotated_image[ + bright_point_1.pixel_y, bright_point_1.pixel_x + ].tolist() == [0, 0, 255] + + +def test_record(sample_image): + tracker = BrightestPointTracker(num_points=2, luminance_threshold=200) + tracked_objects = tracker.process_image(sample_image) + tracker.recorder.record(tracked_objects=tracked_objects) + + assert len(tracker.recorder.recorded_objects) == 1 + + assert len(tracker.recorder.recorded_objects[0]) == 2 diff --git a/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py b/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py index 6acf103..7314538 100644 --- a/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py +++ b/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py @@ -8,13 +8,22 @@ class BrightestPointRecorder(BaseRecorder): def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: - self.recorded_objects.append(deepcopy(tracked_objects["brightest_point"])) + self.recorded_objects.append( + [ + (tracked_object.pixel_x, tracked_object.pixel_y) + for tracked_object in tracked_objects.values() + if "brightest_point" in tracked_object.object_id + ] + ) def process_tracked_objects(self, **kwargs) -> np.ndarray: - self.recorded_objects_array = np.zeros((len(self.recorded_objects), 1, 3)) + num_frames = len(self.recorded_objects) + num_points = len(self.recorded_objects[0]) if num_frames > 0 else 0 + + self.recorded_objects_array = np.zeros((num_frames, num_points, 2)) for i, recorded_object in enumerate(self.recorded_objects): - self.recorded_objects_array[i, 0, 0] = recorded_object.pixel_x - self.recorded_objects_array[i, 0, 1] = recorded_object.pixel_y - self.recorded_objects_array[i, 0, 2] = recorded_object.depth_z + for j, (pixel_x, pixel_y) in enumerate(recorded_object): + self.recorded_objects_array[i, j, 0] = pixel_x + self.recorded_objects_array[i, j, 1] = pixel_y return self.recorded_objects_array diff --git a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py index d934963..094653b 100644 --- a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py +++ b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py @@ -3,10 +3,13 @@ import cv2 import numpy as np +from pydantic import BaseModel from skellytracker.trackers.base_tracker.base_tracker import BaseTracker from skellytracker.trackers.base_tracker.tracked_object import TrackedObject -from skellytracker.trackers.bright_point_tracker.brightest_point_recorder import BrightestPointRecorder +from skellytracker.trackers.bright_point_tracker.brightest_point_recorder import ( + BrightestPointRecorder, +) UPPER_BOUND_COLOR = [255, 255, 255] @@ -15,62 +18,93 @@ logger = logging.getLogger(__name__) +class BrightPatch(BaseModel): + area: float + centroid_x: int + centroid_y: int + + class BrightestPointTracker(BaseTracker): - luminance_threshold: int = 200 + def __init__(self, num_points: int = 1, luminance_threshold: int = 200): + super().__init__( + tracked_object_names=[f"brightest_point_{i}" for i in range(num_points)], + recorder=BrightestPointRecorder(), + ) - def __init__(self): - super().__init__(tracked_object_names=["brightest_point"], recorder=BrightestPointRecorder()) + self.num_points = num_points + self.luminance_threshold = luminance_threshold def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject]: # Convert the image to grayscale gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # Threshold the image to get only bright regions - _, thresholded_image = cv2.threshold(gray_image, self.luminance_threshold, 255, cv2.THRESH_BINARY) + _, thresholded_image = cv2.threshold( + gray_image, self.luminance_threshold, 255, cv2.THRESH_BINARY + ) # Find contours of the bright regions - bright_patches, _ = cv2.findContours(thresholded_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + bright_patches, _ = cv2.findContours( + thresholded_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE + ) # Process each bright patch separately - largest_area = 0 - largest_patch_centroid = None + patch_list = [] for patch in bright_patches: - # Calculate the centroid of the bright patch patch_moments = cv2.moments(patch) if patch_moments["m00"] != 0: # Avoid division by zero centroid_x = int(patch_moments["m10"] / patch_moments["m00"]) centroid_y = int(patch_moments["m01"] / patch_moments["m00"]) - # Keep track of the largest patch and its centroid - area = cv2.contourArea(patch) - if area > largest_area: - largest_area = area - largest_patch_centroid = (centroid_x, centroid_y) + patch_list.append( + BrightPatch( + area=cv2.contourArea(patch), + centroid_x=centroid_x, + centroid_y=centroid_y, + ) + ) + + largest_patches = sorted( + patch_list, key=lambda patch: patch.area, reverse=True + )[: self.num_points] + + for i, patch in enumerate(largest_patches): + self.tracked_objects[f"brightest_point_{i}"].pixel_x = patch.centroid_x + self.tracked_objects[f"brightest_point_{i}"].pixel_y = patch.centroid_y + self.tracked_objects[f"brightest_point_{i}"].extra["thresholdedimage"] = thresholded_image - # If a largest patch was found, update the tracked object - if largest_patch_centroid is not None: - self.tracked_objects["brightest_point"].pixel_x = largest_patch_centroid[0] - self.tracked_objects["brightest_point"].pixel_y = largest_patch_centroid[1] - self.tracked_objects["brightest_point"].extra["thresholded_image"] = thresholded_image + for i in range(len(largest_patches), self.num_points): + self.tracked_objects[f"brightest_point_{i}"].pixel_x = None # TODO: Is this the right value for missing data? + self.tracked_objects[f"brightest_point_{i}"].pixel_y = None - self.annotated_image = self.annotate_image(image=image, - tracked_objects=self.tracked_objects) + self.annotated_image = self.annotate_image( + image=image, tracked_objects=self.tracked_objects + ) return self.tracked_objects - def annotate_image(self, image: np.ndarray, tracked_objects: Dict[str, TrackedObject], **kwargs) -> np.ndarray: - # Copy the original image for annotation + def annotate_image( + self, image: np.ndarray, tracked_objects: Dict[str, TrackedObject], **kwargs + ) -> np.ndarray: annotated_image = image.copy() - # Draw a red 'X' over the largest bright patch - if tracked_objects["brightest_point"].pixel_x is not None and tracked_objects[ - "brightest_point"].pixel_y is not None: - cv2.drawMarker(annotated_image, - (tracked_objects["brightest_point"].pixel_x, tracked_objects["brightest_point"].pixel_y), - (0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2) + for key, tracked_object in tracked_objects.items(): + if ( + "brightest_point" in key + and tracked_object.pixel_x is not None + and tracked_object.pixel_y is not None + ): + cv2.drawMarker( + img=annotated_image, + position=(tracked_object.pixel_x, tracked_object.pixel_y), + color=(0, 0, 255), + markerType=cv2.MARKER_CROSS, + markerSize=20, + thickness=2, + ) return annotated_image if __name__ == "__main__": - BrightestPointTracker().demo() + BrightestPointTracker(num_points=2).demo() From 119c0dd1a0354853d9431803d73b050766086d17 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Mon, 17 Jun 2024 23:11:55 -0600 Subject: [PATCH 15/21] Test YOLO Object Trackers (#40) --- README.md | 31 +++ skellytracker/__init__.py | 2 + .../test_brightest_point_tracker.py | 0 .../tests/test_mediapipe_holistic_tracker.py | 2 - skellytracker/tests/test_test.py | 12 - .../test_yolo_mediapipe_combo_tracker.py | 231 ++++++++++++++++++ .../tests/test_yolo_object_tracker.py | 42 ++++ .../trackers/base_tracker/base_recorder.py | 6 +- .../trackers/base_tracker/base_tracker.py | 12 +- .../charuco_tracker/charuco_tracker.py | 2 +- .../trackers/mmpose_tracker/mmpose_tracker.py | 2 +- .../segment_anything_tracker.py | 4 +- skellytracker/trackers/tracker_manager.py | 44 ---- .../yolo_object_recorder.py | 2 +- .../yolo_object_tracker.py | 8 +- 15 files changed, 327 insertions(+), 73 deletions(-) rename skellytracker/{test => tests}/test_brightest_point_tracker.py (100%) delete mode 100644 skellytracker/tests/test_test.py create mode 100644 skellytracker/tests/test_yolo_mediapipe_combo_tracker.py create mode 100644 skellytracker/tests/test_yolo_object_tracker.py delete mode 100644 skellytracker/trackers/tracker_manager.py diff --git a/README.md b/README.md index cf7b2cd..a8cbbf7 100644 --- a/README.md +++ b/README.md @@ -10,3 +10,34 @@ Then it can be run with `skellytracker`. Running the basic `skellytracker` will open the first webcam port on your computer and run pose estimaiton in realtime with mediapipe holistic as a tracker. You can specify the tracker with `skellytracker TRACKER_NAME`, where `TRACKER_NAME` is the name of an available tracker. To view the names of all available trackers, see `RUN_ME.py`. It will take some time to initialize the tracker the first time you run it, as it will likely need to download the model. + +## Using skellytracker in your project + +To use skellytracker in your project, import a tracker like `from skellytracker import YOLOPoseTracker`, then instantiate it with your desired parameters like `tracker = YOLOPoseTracker(model_size="medium")`, and then use `tracker.process_image(frame)` or `tracker.process_video(video_filepath)`. Processing image by image will let you access each individual annotated frame with `tracker.annotated_image`, and you can optionally record the data with `tracker.recorder.record()`. Access recorded data with `tracker.recorder.process_tracked_objects()`. The running, recording, and processing are done separately to give control over the amount of processing done at each step in the pipeline. Processing an entire video allows you to save the annotated frames as a video, and optionally saves and returns the data as a numpy array. Each tracker has an associated `ModelInfo` class to access model attributes. + +Skellytracker is still under development, so version updates may make breaking changes to the API. Please report any issues and pull requests to the [skellytracker repo](https://github.com/freemocap/skellytracker). + +### Extending the API +To extend the API, import the `BaseTracker` and `BaseRecorder` abstract base classes from skellytracker. Then create a new tracker and recorder inheriting from the base classes and implement all of the abstract methods. + +## Contributing + +We love your input! We want to make contributing to this project as easy and transparent as possible, whether it's: + +- Reporting a bug +- Discussing the current state of the code +- Submitting a fix +- Proposing new features +- Becoming a maintainer + +Pull requests are the best way to propose changes to the codebase (we +use [Github Flow](https://docs.github.com/en/get-started/quickstart/github-flow)). We actively welcome your pull +requests: + +1. Fork the repo and create your branch from `main`. +2. Download the development dependencies with `pip install -e '.[dev]'`. +2. If you've added code that should be tested (including any tracker), add tests. +3. If you've changed APIs, update the documentation. +4. Ensure the test suite passes by running `pytest skellytracker/tests`. +5. Make sure your code lints. +6. Make that pull request! diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index 3b8046c..b0512c8 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -27,10 +27,12 @@ try: from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import MediapipeHolisticTracker + from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import MediapipeModelInfo except: print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker + from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOModelInfo except: print("To use yolo_tracker, install skellytracker[yolo]") try: diff --git a/skellytracker/test/test_brightest_point_tracker.py b/skellytracker/tests/test_brightest_point_tracker.py similarity index 100% rename from skellytracker/test/test_brightest_point_tracker.py rename to skellytracker/tests/test_brightest_point_tracker.py diff --git a/skellytracker/tests/test_mediapipe_holistic_tracker.py b/skellytracker/tests/test_mediapipe_holistic_tracker.py index 58026a5..8c53bdb 100644 --- a/skellytracker/tests/test_mediapipe_holistic_tracker.py +++ b/skellytracker/tests/test_mediapipe_holistic_tracker.py @@ -1,5 +1,3 @@ -import math -import cv2 import pytest import numpy as np diff --git a/skellytracker/tests/test_test.py b/skellytracker/tests/test_test.py deleted file mode 100644 index a3abbfd..0000000 --- a/skellytracker/tests/test_test.py +++ /dev/null @@ -1,12 +0,0 @@ -def returnTrue(num): - try: - return True - except: - return False - - -def test_test(): - """This is a test of the test framework. It should always pass. - To make your own tests, copy this function, change the name, and add your own assertions. - """ - assert returnTrue(6) == True \ No newline at end of file diff --git a/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py b/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py new file mode 100644 index 0000000..72bd8fb --- /dev/null +++ b/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py @@ -0,0 +1,231 @@ +import pytest +import numpy as np + + +from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import ( + MediapipeModelInfo, +) +from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import ( + YOLOMediapipeComboTracker, +) + + +@pytest.mark.usefixtures("test_image") +def test_process_image(test_image): + tracker = YOLOMediapipeComboTracker( + model_size="nano", + model_complexity=0, + ) + tracked_objects = tracker.process_image(test_image) + + assert len(tracked_objects) == 4 + assert tracked_objects["pose_landmarks"] is not None + assert tracked_objects["pose_landmarks"].extra["landmarks"] is not None + assert tracked_objects["right_hand_landmarks"] is not None + assert tracked_objects["right_hand_landmarks"].extra["landmarks"] is not None + assert tracked_objects["left_hand_landmarks"] is not None + assert tracked_objects["left_hand_landmarks"].extra["landmarks"] is not None + assert tracked_objects["face_landmarks"] is not None + assert tracked_objects["face_landmarks"].extra["landmarks"] is not None + + +@pytest.mark.usefixtures("test_image") +def test_annotate_image(test_image): + tracker = YOLOMediapipeComboTracker( + model_size="nano", + model_complexity=0, + ) + tracker.process_image(test_image) + + assert tracker.annotated_image is not None + + +@pytest.mark.usefixtures("test_image") +def test_record_no_buffer(test_image): + tracker = YOLOMediapipeComboTracker( + model_size="nano", + model_complexity=0, + bounding_box_buffer_percentage=0, + ) + tracked_objects = tracker.process_image(test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + assert len(tracker.recorder.recorded_objects[0]) == 4 + + processed_results = tracker.recorder.process_tracked_objects( + image_size=test_image.shape[:2] + ) + assert processed_results is not None + assert processed_results.shape == ( + 1, + MediapipeModelInfo.num_tracked_points_total, + 3, + ) + + expected_results = np.array( + [ + [ + [735.7643890380859, 77.78585314750671, -485.70934295654297], + [757.0420074462891, 73.8272774219513, -451.7356872558594], + [765.8688354492188, 75.42142689228058, -452.0623016357422], + [774.5138549804688, 77.09728181362152, -452.29907989501953], + [729.3473052978516, 70.99358797073364, -452.5171661376953], + [720.2278137207031, 70.5675083398819, -452.57137298583984], + [711.7780303955078, 70.23908793926239, -452.44258880615234], + [780.3971099853516, 86.27676665782928, -254.8146629333496], + [694.5964813232422, 76.93311989307404, -251.65258407592773], + [745.7817077636719, 89.95153248310089, -411.6321563720703], + [709.8857879638672, 86.26414954662323, -411.1183166503906], + [817.0114135742188, 149.35277938842773, -178.73090744018555], + [546.7483139038086, 128.01610708236694, -154.72180366516113], + [825.9496307373047, 219.85299110412598, -142.79363632202148], + [394.10709381103516, 183.29222917556763, -92.53458023071289], + [850.7412719726562, 284.1442108154297, -277.7159309387207], + [251.08980178833008, 229.20471668243408, -194.44988250732422], + [860.8124542236328, 301.0432004928589, -309.02509689331055], + [211.39860153198242, 242.8161120414734, -219.8578643798828], + [843.7083435058594, 300.99024295806885, -385.36643981933594], + [222.63912200927734, 246.6996932029724, -290.17900466918945], + [839.26025390625, 295.74596643447876, -312.1304130554199], + [240.05075454711914, 241.5028166770935, -226.83252334594727], + [632.852668762207, 273.6677813529968, -22.331013679504395], + [484.76829528808594, 259.1326332092285, 21.80124521255493], + [549.3550491333008, 360.4587650299072, -38.61574411392212], + [425.41954040527344, 343.2985496520996, -20.899696350097656], + [486.68445587158203, 415.20827293395996, 411.0944366455078], + [399.84127044677734, 408.2763719558716, 398.0601501464844], + [476.32495880126953, 416.31832122802734, 451.4645767211914], + ] + ] + ) + assert np.allclose( + processed_results[:, :30, :], expected_results[:, :30, :], atol=2 + ) + + +@pytest.mark.usefixtures("test_image") +def test_record_buffer_by_image_size(test_image): + tracker = YOLOMediapipeComboTracker( + model_size="nano", + model_complexity=0, + bounding_box_buffer_percentage=10, + buffer_size_method="buffer_by_image_size", + ) + tracked_objects = tracker.process_image(test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + assert len(tracker.recorder.recorded_objects[0]) == 4 + + processed_results = tracker.recorder.process_tracked_objects( + image_size=test_image.shape[:2] + ) + assert processed_results is not None + assert processed_results.shape == ( + 1, + MediapipeModelInfo.num_tracked_points_total, + 3, + ) + + expected_results = np.array( + [ + [ + [732.0687866210938, 79.0345823764801, -635.7109069824219], + [753.8626098632812, 75.13578772544861, -608.0960464477539], + [761.7510223388672, 76.88950717449188, -608.355827331543], + [769.5530700683594, 78.81806373596191, -608.4667587280273], + [729.2195892333984, 72.18442976474762, -607.5491333007812], + [721.5387725830078, 71.86003804206848, -607.6824188232422], + [714.1613006591797, 71.62245333194733, -607.6252746582031], + [775.9226226806641, 86.93966388702393, -409.65084075927734], + [698.8643646240234, 77.10046291351318, -402.1929931640625], + [741.6652679443359, 90.84179520606995, -557.6795196533203], + [709.4770050048828, 86.31318032741547, -555.5589294433594], + [814.6518707275391, 149.26713109016418, -293.33160400390625], + [544.8312759399414, 127.3474645614624, -270.33700942993164], + [831.3512420654297, 219.2579483985901, -218.43671798706055], + [392.59105682373047, 186.1086130142212, -182.6705551147461], + [854.2007446289062, 282.19802141189575, -344.9155044555664], + [248.35845947265625, 233.24616193771362, -273.1560516357422], + [857.7934265136719, 299.5019817352295, -382.82962799072266], + [202.44756698608398, 248.4213924407959, -300.88146209716797], + [844.1407775878906, 300.61514139175415, -460.82714080810547], + [212.99942016601562, 251.683087348938, -378.4083938598633], + [839.8464965820312, 294.21818017959595, -380.2419662475586], + [230.61071395874023, 246.78754091262817, -307.7861785888672], + [639.9309539794922, 268.61598014831543, -0.536465011537075], + [481.90975189208984, 254.10432815551758, 0.47791849821805954], + [546.6202926635742, 351.9411635398865, 89.65752601623535], + [416.8575668334961, 342.4706482887268, -87.49273300170898], + [496.87782287597656, 406.89836025238037, 635.2282333374023], + [393.65081787109375, 407.12096214294434, 370.8687210083008], + [481.2765884399414, 408.57420444488525, 681.8212890625], + ] + ] + ) + assert np.allclose( + processed_results[:, :30, :], expected_results[:, :30, :], atol=2 + ) + + +@pytest.mark.usefixtures("test_image") +def test_record_buffer_by_box_size(test_image): + tracker = YOLOMediapipeComboTracker( + model_size="nano", + model_complexity=0, + bounding_box_buffer_percentage=10, + buffer_size_method="buffer_by_box_size", + ) + tracked_objects = tracker.process_image(test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + assert len(tracker.recorder.recorded_objects[0]) == 4 + + processed_results = tracker.recorder.process_tracked_objects( + image_size=test_image.shape[:2] + ) + assert processed_results is not None + assert processed_results.shape == ( + 1, + MediapipeModelInfo.num_tracked_points_total, + 3, + ) + + expected_results = np.array( + [ + [ + [731.2718200683594, 77.88420975208282, -548.9945602416992], + [754.6127319335938, 74.30741965770721, -521.489372253418], + [762.8125762939453, 76.33775532245636, -521.7532348632812], + [771.2681579589844, 78.52324604988098, -521.8650436401367], + [730.0675964355469, 70.9510749578476, -521.0283660888672], + [722.1942138671875, 70.54689288139343, -521.1295318603516], + [714.8076629638672, 70.24498343467712, -521.0332870483398], + [779.2241668701172, 87.20496118068695, -340.54332733154297], + [700.3346252441406, 76.72817766666412, -330.6981658935547], + [741.3204956054688, 90.19120931625366, -479.1018295288086], + [708.2732391357422, 85.63711881637573, -476.2978744506836], + [813.9542388916016, 149.21152353286743, -249.2682647705078], + [545.847282409668, 128.2720971107483, -214.6741485595703], + [831.6414642333984, 217.95576810836792, -199.20928955078125], + [392.5309371948242, 185.60349941253662, -141.00683212280273], + [850.6895446777344, 280.559663772583, -331.45111083984375], + [252.97996520996094, 230.9888792037964, -229.0725326538086], + [856.6841125488281, 298.4851026535034, -367.7482604980469], + [203.86322021484375, 247.05041885375977, -250.86517333984375], + [841.0806274414062, 299.39956426620483, -436.93775177001953], + [215.9295654296875, 250.416419506073, -324.9034881591797], + [836.7790985107422, 292.6754379272461, -363.13419342041016], + [234.16423797607422, 246.12889766693115, -261.9901657104492], + [639.0184783935547, 271.9633984565735, -11.131852865219116], + [485.26206970214844, 257.6048684120178, 10.823948383331299], + [544.7049331665039, 353.29198837280273, 29.705591201782227], + [422.3688507080078, 341.57193660736084, -53.83963108062744], + [484.64855194091797, 412.5028896331787, 527.0954132080078], + [392.8578186035156, 406.18433475494385, 385.8837127685547], + [471.3254165649414, 411.8346977233887, 569.1376495361328], + ] + ] + ) + assert np.allclose( + processed_results[:, :30, :], expected_results[:, :30, :], atol=2 + ) diff --git a/skellytracker/tests/test_yolo_object_tracker.py b/skellytracker/tests/test_yolo_object_tracker.py new file mode 100644 index 0000000..4cd0843 --- /dev/null +++ b/skellytracker/tests/test_yolo_object_tracker.py @@ -0,0 +1,42 @@ +import pytest +import numpy as np + + +from skellytracker.trackers.yolo_object_tracker.yolo_object_tracker import ( + YOLOObjectTracker, +) + + +@pytest.mark.usefixtures("test_image") +def test_process_image_person_only(test_image): + tracker = YOLOObjectTracker(model_size="nano", person_only=True) + tracked_objects = tracker.process_image(test_image) + + assert len(tracked_objects) == 1 + assert tracked_objects["object"] is not None + assert tracked_objects["object"].extra["boxes_xyxy"] is not None + assert np.allclose(tracked_objects["object"].extra["boxes_xyxy"], [90.676,96.981,493.54,812.03], atol=1e-2) + assert tracked_objects["object"].extra["original_image_shape"] == (1280, 720) + +@pytest.mark.usefixtures("test_image") +def test_annotate_image(test_image): + tracker = YOLOObjectTracker() + tracker.process_image(test_image) + + assert tracker.annotated_image is not None + + +@pytest.mark.usefixtures("test_image") +def test_record(test_image): + tracker = YOLOObjectTracker(model_size="nano", person_only=True) + tracked_objects = tracker.process_image(test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + + processed_results = tracker.recorder.process_tracked_objects() + assert processed_results is not None + assert processed_results.shape == (1,4) + + assert np.allclose( + processed_results, [90.676,96.981,493.54,812.03], atol=1e-2 + ) diff --git a/skellytracker/trackers/base_tracker/base_recorder.py b/skellytracker/trackers/base_tracker/base_recorder.py index 9974da0..9a0fc6e 100644 --- a/skellytracker/trackers/base_tracker/base_recorder.py +++ b/skellytracker/trackers/base_tracker/base_recorder.py @@ -53,6 +53,8 @@ def save(self, file_path: str) -> None: :return: None """ if self.recorded_objects_array is None: - self.process_tracked_objects() + recorded_objects_array = self.process_tracked_objects() + else: + recorded_objects_array = self.recorded_objects_array logger.info(f"Saving recorded objects to {file_path}") - np.save(file_path, self.recorded_objects_array) + np.save(file_path, recorded_objects_array) diff --git a/skellytracker/trackers/base_tracker/base_tracker.py b/skellytracker/trackers/base_tracker/base_tracker.py index a3ab1c5..a21c0ed 100644 --- a/skellytracker/trackers/base_tracker/base_tracker.py +++ b/skellytracker/trackers/base_tracker/base_tracker.py @@ -25,8 +25,8 @@ class BaseTracker(ABC): def __init__( self, - tracked_object_names: List[str] = None, - recorder: BaseRecorder = None, + recorder: BaseRecorder, + tracked_object_names: List[str] = [], **data: Any, ): self.recorder = recorder @@ -65,7 +65,7 @@ def process_video( output_video_filepath: Optional[Union[str, Path]] = None, save_data_bool: bool = False, use_tqdm: bool = True, - ) -> np.ndarray: + ) -> Optional[np.ndarray]: """ Run the tracker on a video. @@ -73,7 +73,7 @@ def process_video( :param output_video_filepath: Path to save annotated video to, does not save video if None. :param save_data_bool: Whether to save the data to a file. :param use_tqdm: Whether to use tqdm to show a progress bar - :return: Array of tracked keypoint data + :return: Array of tracked keypoint data, if save_data_bool is True """ cap = cv2.VideoCapture(str(input_video_filepath)) @@ -118,6 +118,8 @@ def process_video( if self.recorder is not None: self.recorder.record(self.tracked_objects) if video_handler is not None: + if self.annotated_image is None: + self.annotated_image = frame video_handler.add_frame(self.annotated_image) ret, frame = cap.read() @@ -130,7 +132,7 @@ def process_video( output_array = self.recorder.process_tracked_objects(image_size=image_size) if save_data_bool: self.recorder.save( - file_path=Path(input_video_filepath).with_suffix(".npy") + file_path=str(Path(input_video_filepath).with_suffix(".npy")) ) else: output_array = None diff --git a/skellytracker/trackers/charuco_tracker/charuco_tracker.py b/skellytracker/trackers/charuco_tracker/charuco_tracker.py index c0d172c..46c72a3 100644 --- a/skellytracker/trackers/charuco_tracker/charuco_tracker.py +++ b/skellytracker/trackers/charuco_tracker/charuco_tracker.py @@ -16,7 +16,7 @@ def __init__(self, squareLength: float = 1, markerLength: float = .8, ): - super().__init__(tracked_object_names=tracked_object_names) + super().__init__(recorder=None, tracked_object_names=tracked_object_names) self.board = cv2.aruco.CharucoBoard_create(squares_x, squares_y, squareLength, markerLength, dictionary) def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject]: diff --git a/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py b/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py index 64e6b62..abcf4c1 100644 --- a/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py +++ b/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py @@ -7,7 +7,7 @@ class MMPoseTracker(BaseTracker): def __init__(self, config_file, checkpoint_file): - super().__init__(tracked_object_names=["human_pose"]) + super().__init__(recorder=None, tracked_object_names=["human_pose"]) self.model = init_pose_model(config_file, checkpoint_file, device='cuda:0') def process_image(self, image, **kwargs): diff --git a/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py b/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py index 4cd327e..f9a8882 100644 --- a/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py +++ b/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py @@ -6,8 +6,8 @@ from skellytracker.trackers.base_tracker.base_tracker import BaseTracker class SAMTracker(BaseTracker): - def __init__(self, model_size: str="nano"): - super().__init__(tracked_object_names=["segmentation"]) + def __init__(self): + super().__init__(recorder=None,tracked_object_names=["segmentation"]) self.model = SAM('sam_b.pt') diff --git a/skellytracker/trackers/tracker_manager.py b/skellytracker/trackers/tracker_manager.py deleted file mode 100644 index 63f7af5..0000000 --- a/skellytracker/trackers/tracker_manager.py +++ /dev/null @@ -1,44 +0,0 @@ -import multiprocessing as mp -import time -from asyncio import sleep -from typing import List - -from skellytracker.trackers.base_tracker.base_tracker import BaseTracker -from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import BrightestPointTracker - - -class TrackerManager: - def __init__(self, trackers: List[BaseTracker]): - self.trackers = trackers - self.parent_connection, self.child_connection = mp.Pipe() - self.process = mp.Process(target=self._process_images, args=(self.child_connection, self.trackers)) - self.process.start() - - @staticmethod - def _process_images(conn, trackers): - while True: - time.sleep(0.001) - image = conn.recv() - if image is None: - break - for tracker in trackers: - tracker.process_image(image) - - def add_image(self, image): - self.parent_connection.send(image) - - def demo(self): - self.trackers[0].demo() - - def stop(self): - self.parent_connection.send(None) - self.process.join() - - - -if __name__ == "__main__": - - trackers = [BrightestPointTracker()] - - manager = TrackerManager(trackers) - manager.demo() diff --git a/skellytracker/trackers/yolo_object_tracker/yolo_object_recorder.py b/skellytracker/trackers/yolo_object_tracker/yolo_object_recorder.py index 5f785e9..3afea7e 100644 --- a/skellytracker/trackers/yolo_object_tracker/yolo_object_recorder.py +++ b/skellytracker/trackers/yolo_object_tracker/yolo_object_recorder.py @@ -13,6 +13,6 @@ def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: def process_tracked_objects(self, **kwargs) -> np.ndarray: self.recorded_objects_array = np.zeros((len(self.recorded_objects), 4)) for i, recorded_object in enumerate(self.recorded_objects): - self.recorded_objects_array[i, :] = recorded_object.extra["boxes_xywh"] + self.recorded_objects_array[i, :] = recorded_object.extra["boxes_xyxy"] return self.recorded_objects_array diff --git a/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py b/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py index 62dca7b..0257bbf 100644 --- a/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py +++ b/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py @@ -24,21 +24,23 @@ def __init__( pytorch_model = yolo_object_model_dictionary[model_size] self.model = YOLO(pytorch_model) self.confidence_threshold = confidence_threshold + # TODO: When we expose this in freemocap, replace this with an int/list[int] to decide which class to track + # TODO: Will also need to parameterize the "max_det" and setup tracker to take multiple tracked objects if person_only: self.classes = 0 # 0 is the YOLO class for person detection else: - self.classes = None + self.classes = None # None includes all classes def process_image(self, image, **kwargs) -> Dict[str, TrackedObject]: results = self.model(image, classes=self.classes, max_det=1, verbose=False, conf=self.confidence_threshold) - box_xyxy = np.asarray(results[0].boxes.xyxy.cpu()).flatten() # if on GPU, need to copy to CPU before np array conversion + box_xyxy = np.asarray(results[0].boxes.xyxy.cpu()).flatten() # On GPU, need to copy to CPU before np array conversion if box_xyxy.size > 0: self.tracked_objects["object"].pixel_x = (box_xyxy[0] + box_xyxy[2]) / 0.5 self.tracked_objects["object"].pixel_y = (box_xyxy[1] + box_xyxy[3]) / 0.5 - self.tracked_objects["object"].extra["boxes_xywy"] = box_xyxy + self.tracked_objects["object"].extra["boxes_xyxy"] = box_xyxy self.tracked_objects["object"].extra["original_image_shape"] = results[ 0 ].boxes.orig_shape From 54058ff1959d6f0a50fbcac9ea44f45875b10cec Mon Sep 17 00:00:00 2001 From: philipqueen Date: Tue, 25 Jun 2024 10:43:39 -0600 Subject: [PATCH 16/21] Philip/support 3.12 (#36) --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e4bd6df..711e06f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,15 +60,15 @@ dependencies = [ "pydantic==1.*", "tqdm==4.*", ] -requires-python = ">=3.9,<3.12" +requires-python = ">=3.9,<3.13" dynamic = ["version", "description"] [project.optional-dependencies] dev = ["black", "bumpver", "isort", "pip-tools", "pytest"] -mediapipe = ["mediapipe==0.10.9"] +mediapipe = ["mediapipe==0.10.14"] yolo = ["ultralytics~=8.0.202"] -all = ["ultralytics~=8.0.202", "mediapipe==0.10.9"] +all = ["ultralytics~=8.0.202", "mediapipe==0.10.14"] [project.urls] Homepage = "https://github.com/freemocap/skellytracker" From 9d1472f9c6250faf2a95eb612e5015768ea99948 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Tue, 25 Jun 2024 10:44:01 -0600 Subject: [PATCH 17/21] Pydantic 2 upgrade (#35) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 711e06f..c780555 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,8 +56,8 @@ keywords = [ #dynamic = ["dependencies"] dependencies = [ "opencv-contrib-python==4.8.*", + "pydantic==2.*", "numpy<2", - "pydantic==1.*", "tqdm==4.*", ] requires-python = ">=3.9,<3.13" From d5edf7fad568a61601402c6b3bed7ae83f3800f3 Mon Sep 17 00:00:00 2001 From: jonmatthis Date: Tue, 25 Jun 2024 12:46:08 -0400 Subject: [PATCH 18/21] Bump version v2024.06.1015 -> v2024.06.1016 --- pyproject.toml | 2 +- skellytracker/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c780555..85e11e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,7 +74,7 @@ all = ["ultralytics~=8.0.202", "mediapipe==0.10.14"] Homepage = "https://github.com/freemocap/skellytracker" [tool.bumpver] -current_version = "v2024.06.1015" +current_version = "v2024.06.1016" version_pattern = "vYYYY.0M.BUILD[-TAG]" commit_message = "Bump version {old_version} -> {new_version}" diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index b0512c8..ab6a855 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -1,7 +1,7 @@ """Top-level package for skellytracker""" __package_name__ = "skellytracker" -__version__ = "v2024.06.1015" +__version__ = "v2024.06.1016" __author__ = """Skelly FreeMoCap""" __email__ = "info@freemocap.org" From 447e95cb3c23a3d07e02ea92f5a5491bd2596f3d Mon Sep 17 00:00:00 2001 From: philipqueen Date: Thu, 4 Jul 2024 11:27:37 -0600 Subject: [PATCH 19/21] Revamp charuco tracker (#41) --- skellytracker/RUN_ME.py | 10 +- skellytracker/SINGLE_IMAGE_RUN.py | 10 +- skellytracker/__init__.py | 2 - skellytracker/system/default_paths.py | 1 + skellytracker/tests/conftest.py | 9 +- skellytracker/tests/test_charuco_tracker.py | 148 ++++++++++++++++++ .../brightest_point_tracker.py | 2 +- .../charuco_tracker/charuco_recorder.py | 23 +++ .../charuco_tracker/charuco_tracker.py | 133 ++++++++++------ .../mediapipe_holistic_recorder.py | 2 +- .../trackers/yolo_tracker/yolo_recorder.py | 2 +- 11 files changed, 288 insertions(+), 54 deletions(-) create mode 100644 skellytracker/tests/test_charuco_tracker.py create mode 100644 skellytracker/trackers/charuco_tracker/charuco_recorder.py diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index 5a04c8b..e09153e 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -29,9 +29,15 @@ def main(demo_tracker: str = "mediapipe_holistic_tracker"): BrightestPointTracker(num_points=2).demo() elif demo_tracker == "charuco_tracker": + charuco_squares_x_in = 7 + charuco_squares_y_in = 5 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + CharucoTracker( - squaresX=7, - squaresY=5, + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), ).demo() diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index f1d337d..02f112b 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -26,9 +26,15 @@ BrightestPointTracker(num_points=2).image_demo(image_path=image_path) elif demo_tracker == "charuco_tracker": + charuco_squares_x_in = 7 + charuco_squares_y_in = 5 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + CharucoTracker( - squaresX=7, - squaresY=5, + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), ).image_demo(image_path=image_path) diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index ab6a855..91413a8 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -20,8 +20,6 @@ print(f"adding base_package_path: {base_package_path} : to sys.path") sys.path.insert(0, str(base_package_path)) # add parent directory to sys.path -print(f"sys path: {sys.path}") - from skellytracker.system.default_paths import get_log_file_path from skellytracker.system.logging_configuration import configure_logging diff --git a/skellytracker/system/default_paths.py b/skellytracker/system/default_paths.py index 58a5d4c..a570c60 100644 --- a/skellytracker/system/default_paths.py +++ b/skellytracker/system/default_paths.py @@ -8,6 +8,7 @@ LOGS_INFO_AND_SETTINGS_FOLDER_NAME = "logs_info_and_settings" LOG_FILE_FOLDER_NAME = "logs" FIGSHARE_TEST_IMAGE_URL = "https://figshare.com/ndownloader/files/47043898" +FIGSHARE_CHARUCO_TEST_IMAGE_URL = "https://figshare.com/ndownloader/files/47127685" def get_base_folder_path(): base_folder = Path().home() / BASE_FOLDER_NAME diff --git a/skellytracker/tests/conftest.py b/skellytracker/tests/conftest.py index 3f226dd..3c4ca48 100644 --- a/skellytracker/tests/conftest.py +++ b/skellytracker/tests/conftest.py @@ -1,14 +1,21 @@ import numpy as np import pytest +from skellytracker.system.default_paths import FIGSHARE_CHARUCO_TEST_IMAGE_URL from skellytracker.utilities.download_test_image import download_test_image class SessionInfo: test_image: np.ndarray + charuco_test_image: np.ndarray def pytest_sessionstart(session): SessionInfo.test_image = download_test_image() + SessionInfo.charuco_test_image = download_test_image(test_image_url=FIGSHARE_CHARUCO_TEST_IMAGE_URL) @pytest.fixture def test_image(): - return SessionInfo.test_image \ No newline at end of file + return SessionInfo.test_image + +@pytest.fixture +def charuco_test_image(): + return SessionInfo.charuco_test_image \ No newline at end of file diff --git a/skellytracker/tests/test_charuco_tracker.py b/skellytracker/tests/test_charuco_tracker.py new file mode 100644 index 0000000..235b77c --- /dev/null +++ b/skellytracker/tests/test_charuco_tracker.py @@ -0,0 +1,148 @@ +import math +import cv2 +import pytest +import numpy as np + + +from skellytracker.trackers.charuco_tracker.charuco_tracker import CharucoTracker + + +@pytest.mark.usefixtures("charuco_test_image") +def test_process_image(charuco_test_image): + charuco_squares_x_in = 7 + charuco_squares_y_in = 5 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + + tracker = CharucoTracker( + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ) + tracked_objects = tracker.process_image(charuco_test_image) + + expected_results = { + "0": (307.99796, 110.00571), + "1": (336.93832, 120.001335), + "2": (366.70923, 130.58067), + "3": (396.36816, 141.78345), + "4": (425.60236, 153.44989), + "5": (455.322, 165.67236), + "6": (294.39023, 135.20029), + "7": (323.91107, 145.17943), + "8": (353.4799, 155.70189), + "9": (383.16925, 167.31921), + "10": (412.6318, 179.24583), + "11": (442.47086, 191.43738), + "12": (280.9244, 160.91164), + "13": (310.19556, 171.59216), + "14": (339.86594, 182.31856), + "15": (369.6749, 193.77588), + "16": (399.5786, 205.84789), + "17": (429.53903, 218.09035), + "18": (267.12946, 187.75053), + "19": (296.7608, 198.35608), + "20": (326.31488, 209.67616), + "21": (356.24832, 220.93948), + "22": (386.3587, 233.16882), + "23": (416.3993, 245.57997), + } + + assert len(tracked_objects) == len(charuco_ids) + for id in charuco_ids: + tracked_corner = tracked_objects[id] + assert tracked_corner.pixel_x is not None + assert tracked_corner.pixel_y is not None + assert np.allclose( + (tracked_corner.pixel_x, tracked_corner.pixel_y), expected_results[id] + ) + + +@pytest.mark.usefixtures("test_image") +def test_image_without_charuco(test_image): + charuco_squares_x_in = 7 + charuco_squares_y_in = 5 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + + tracker = CharucoTracker( + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ) + tracked_objects = tracker.process_image(test_image) + + assert len(tracked_objects) == len(charuco_ids) + for id in charuco_ids: + tracked_corner = tracked_objects[id] + assert tracked_corner.pixel_x is None + assert tracked_corner.pixel_y is None + + +@pytest.mark.usefixtures("charuco_test_image") +def test_annotate_image(charuco_test_image): + charuco_squares_x_in = 7 + charuco_squares_y_in = 5 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + + tracker = CharucoTracker( + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ) + tracker.process_image(charuco_test_image) + + assert tracker.annotated_image is not None + assert not np.all(tracker.annotated_image == charuco_test_image) + + +@pytest.mark.usefixtures("charuco_test_image") +def test_record(charuco_test_image): + charuco_squares_x_in = 7 + charuco_squares_y_in = 5 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + + tracker = CharucoTracker( + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ) + tracked_objects = tracker.process_image(charuco_test_image) + tracker.recorder.record(tracked_objects=tracked_objects) + assert len(tracker.recorder.recorded_objects) == 1 + + processed_results = tracker.recorder.process_tracked_objects() + assert processed_results is not None + assert processed_results.shape == (1, len(charuco_ids), 2) + + # expected_results = np.array( + # [ + # [ + # [392.56927490234375, 140.40118408203125, np.nan], + # [414.94940185546875, 132.90655517578125, np.nan], + # [386.4353332519531, 125.51483154296875, np.nan], + # [446.2061767578125, 157.98883056640625, np.nan], + # [373.6619873046875, 138.93646240234375, np.nan], + # [453.78662109375, 265.081787109375, np.nan], + # [317.9375305175781, 231.9653778076172, np.nan], + # [465.893310546875, 396.12274169921875, np.nan], + # [220.12176513671875, 325.96636962890625, np.nan], + # [465.23358154296875, 499.0487365722656, np.nan], + # [142.17066955566406, 407.7397155761719, np.nan], + # [352.325439453125, 468.44671630859375, np.nan], + # [268.8867492675781, 448.7227783203125, np.nan], + # [310.899658203125, 630.5478515625, np.nan], + # [227.7810821533203, 617.9011840820312, np.nan], + # [269.08587646484375, 733.4285888671875, np.nan], + # [213.1557159423828, 741.54541015625, np.nan], + # ] + # ] + # ) + # assert np.allclose(processed_results[:, :, :2], expected_results[:, :, :2], atol=1e-2) + # assert np.isnan(processed_results[:, :, 2]).all() diff --git a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py index 094653b..a5f529d 100644 --- a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py +++ b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py @@ -96,7 +96,7 @@ def annotate_image( ): cv2.drawMarker( img=annotated_image, - position=(tracked_object.pixel_x, tracked_object.pixel_y), + position=(int(tracked_object.pixel_x), int(tracked_object.pixel_y)), color=(0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=20, diff --git a/skellytracker/trackers/charuco_tracker/charuco_recorder.py b/skellytracker/trackers/charuco_tracker/charuco_recorder.py new file mode 100644 index 0000000..b048f85 --- /dev/null +++ b/skellytracker/trackers/charuco_tracker/charuco_recorder.py @@ -0,0 +1,23 @@ +from copy import deepcopy +from typing import Dict +import numpy as np + +from skellytracker.trackers.base_tracker.base_recorder import BaseRecorder +from skellytracker.trackers.base_tracker.tracked_object import TrackedObject + + +class CharucoRecorder(BaseRecorder): + def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: + self.recorded_objects.append( + [deepcopy(tracked_object) for tracked_object in tracked_objects.values()] + ) + + def process_tracked_objects(self, **kwargs) -> np.ndarray: + self.recorded_objects_array = np.array( + [ + [[tracked_object.pixel_x, tracked_object.pixel_y] for tracked_object in tracked_object_list] + for tracked_object_list in self.recorded_objects + ] + ) + + return self.recorded_objects_array diff --git a/skellytracker/trackers/charuco_tracker/charuco_tracker.py b/skellytracker/trackers/charuco_tracker/charuco_tracker.py index 46c72a3..62d7435 100644 --- a/skellytracker/trackers/charuco_tracker/charuco_tracker.py +++ b/skellytracker/trackers/charuco_tracker/charuco_tracker.py @@ -5,72 +5,117 @@ from skellytracker.trackers.base_tracker.base_tracker import BaseTracker from skellytracker.trackers.base_tracker.tracked_object import TrackedObject +from skellytracker.trackers.charuco_tracker.charuco_recorder import CharucoRecorder class CharucoTracker(BaseTracker): - def __init__(self, - tracked_object_names: List[str], - squares_x: int, - squares_y: int, - dictionary: cv2.aruco_Dictionary, - squareLength: float = 1, - markerLength: float = .8, - ): - super().__init__(recorder=None, tracked_object_names=tracked_object_names) - self.board = cv2.aruco.CharucoBoard_create(squares_x, squares_y, squareLength, markerLength, dictionary) + def __init__( + self, + tracked_object_names: List[str], + squares_x: int, + squares_y: int, + dictionary: cv2.aruco.Dictionary = cv2.aruco.getPredefinedDictionary( + cv2.aruco.DICT_4X4_250 + ), + square_length: float = 1, + marker_length: float = 0.8, + ): + super().__init__( + recorder=CharucoRecorder(), tracked_object_names=tracked_object_names + ) + self.board = cv2.aruco.CharucoBoard( + size=(squares_x, squares_y), + squareLength=square_length, + markerLength=marker_length, + dictionary=dictionary, + ) + + # Following most recent charuco detection documentation: https://docs.opencv.org/4.x/df/d4a/tutorial_charuco_detection.html + self.charuco_detector = cv2.aruco.CharucoDetector(self.board) + + self.tracked_object_names = tracked_object_names + self.dictionary = dictionary def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject]: # Convert the image to grayscale gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - # Detect Aruco markers - corners, ids, _ = cv2.aruco.detectMarkers(gray_image, self.board.dictionary) - - # If any markers were found - if len(corners) > 0: - # Refine the detected markers - ret, charuco_corners, charuco_ids = cv2.aruco.interpolateCornersCharuco(corners, ids, gray_image, - self.board) - - # If any Charuco corners were found - if charuco_corners is not None and charuco_ids is not None and len(charuco_corners) > 3: - # Clear previous tracked objects - self.tracked_objects.clear() - - # Create a TrackedObject for each corner - for i, corner in enumerate(charuco_corners): - object_id = str(i) - self.tracked_objects[object_id] = TrackedObject(object_id=object_id) - self.tracked_objects[object_id].pixel_x = corner[0][0] - self.tracked_objects[object_id].pixel_y = corner[0][1] - - self.annotated_image = self.annotate_image(image=image, - tracked_objects=self.tracked_objects) + charuco_corners, charuco_ids, _marker_corners, _marker_ids = ( + self.charuco_detector.detectBoard(gray_image) + ) + + self.reinitialize_tracked_objects() + + # If any Charuco corners were found + if ( + charuco_corners is not None + and charuco_ids is not None + and len(charuco_corners) > 3 + ): + # Create a TrackedObject for each corner + for id, corner in zip(charuco_ids, charuco_corners): + object_id = str(id).strip("[]") + self.tracked_objects[object_id] = TrackedObject(object_id=object_id) + self.tracked_objects[object_id].pixel_x = corner[0][0] + self.tracked_objects[object_id].pixel_y = corner[0][1] + + self.annotated_image = self.annotate_image( + image=image, tracked_objects=self.tracked_objects + ) return self.tracked_objects - def annotate_image(self, image: np.ndarray, tracked_objects: Dict[str, TrackedObject], **kwargs) -> np.ndarray: + def annotate_image( + self, image: np.ndarray, tracked_objects: Dict[str, TrackedObject], **kwargs + ) -> np.ndarray: # Copy the original image for annotation annotated_image = image.copy() # Draw a marker for each tracked corner for tracked_object in tracked_objects.values(): - if tracked_object.pixel_x is not None and tracked_object.pixel_y is not None: - cv2.drawMarker(annotated_image, - (int(tracked_object.pixel_x), int(tracked_object.pixel_y)), - (0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2) + if ( + tracked_object.pixel_x is not None + and tracked_object.pixel_y is not None + ): + cv2.drawMarker( + annotated_image, + (int(tracked_object.pixel_x), int(tracked_object.pixel_y)), + (0, 0, 255), + markerType=cv2.MARKER_CROSS, + markerSize=30, + thickness=2, + ) + cv2.putText( + annotated_image, + tracked_object.object_id, + (int(tracked_object.pixel_x), int(tracked_object.pixel_y)), + cv2.FONT_HERSHEY_SIMPLEX, + 1, + (255, 0, 0), + 2, + ) return annotated_image + def reinitialize_tracked_objects(self) -> None: + """ + Reinitialize tracked objects to clear previous frames data + + Unlike self.tracked_objects.clear(), this will ensure every tracked object has a value for each frame, even if its empty + """ + for name in self.tracked_object_names: + self.tracked_objects[name] = TrackedObject(object_id=name) + if __name__ == "__main__": charuco_squares_x_in = 7 charuco_squares_y_in = 5 - number_of_charuco_markers = charuco_squares_x_in - 1 * charuco_squares_y_in - 1 + number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) charuco_ids = [str(index) for index in range(number_of_charuco_markers)] - CharucoTracker(tracked_object_names=charuco_ids, - squares_x=charuco_squares_x_in, - squares_y=charuco_squares_y_in, - dictionary=cv2.aruco.Dictionary_get(cv2.aruco.DICT_4X4_250) - ).demo() + CharucoTracker( + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x_in, + squares_y=charuco_squares_y_in, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ).demo() diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py index 9e0f298..de8c9cb 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py @@ -55,7 +55,7 @@ def process_tracked_objects(self, **kwargs) -> np.ndarray: else: number = MediapipeModelInfo.num_tracked_points_right_hand for _ in range(number): - self.recorded_objects_array[i, landmark_number, :] = np.NaN + self.recorded_objects_array[i, landmark_number, :] = np.nan landmark_number += 1 return self.recorded_objects_array diff --git a/skellytracker/trackers/yolo_tracker/yolo_recorder.py b/skellytracker/trackers/yolo_tracker/yolo_recorder.py index dbf4db8..30f2230 100644 --- a/skellytracker/trackers/yolo_tracker/yolo_recorder.py +++ b/skellytracker/trackers/yolo_tracker/yolo_recorder.py @@ -24,6 +24,6 @@ def process_tracked_objects(self, **kwargs) -> np.ndarray: self.recorded_objects_array[i, j, 1] = recorded_object.extra[ "landmarks" ][0, j, 1] - self.recorded_objects_array[i, j, 2] = np.NaN + self.recorded_objects_array[i, j, 2] = np.nan return self.recorded_objects_array From d3f9cd4c1cfa9958cbe6d7afc0a73561454c3818 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Tue, 30 Jul 2024 12:35:58 -0600 Subject: [PATCH 20/21] Linting and type checking (#44) --- .github/workflows/linting-ruff.yml | 27 ++++++++++++ pyproject.toml | 8 +++- skellytracker/RUN_ME.py | 21 ++++----- skellytracker/SINGLE_IMAGE_RUN.py | 15 +++---- skellytracker/__init__.py | 31 ++++++++----- skellytracker/__main__.py | 11 ++--- skellytracker/process_folder_of_videos.py | 19 ++++---- skellytracker/system/default_paths.py | 19 ++++++-- skellytracker/system/logging_configuration.py | 3 +- skellytracker/tests/conftest.py | 6 ++- .../tests/test_brightest_point_tracker.py | 11 +++-- skellytracker/tests/test_charuco_tracker.py | 1 - .../tests/test_yolo_object_tracker.py | 13 +++--- skellytracker/tests/test_yolo_pose_tracker.py | 10 ++++- .../trackers/base_tracker/base_tracker.py | 8 ++-- .../trackers/base_tracker/tracked_object.py | 3 +- .../trackers/base_tracker/video_handler.py | 6 +-- .../brightest_point_recorder.py | 1 - .../brightest_point_tracker.py | 8 +++- .../charuco_tracker/charuco_tracker.py | 6 +-- .../image_demo_viewer.py | 6 ++- .../webcam_demo_viewer.py | 3 +- .../mediapipe_holistic_recorder.py | 4 +- .../trackers/mmpose_tracker/mmpose_tracker.py | 23 ++++------ .../segment_anything_tracker.py | 13 +++--- .../yolo_object_tracker.py | 12 +++++- .../trackers/yolo_tracker/yolo_tracker.py | 12 ++++-- .../utilities/download_test_image.py | 6 +-- skellytracker/utilities/get_video_paths.py | 2 +- .../quine_directory_printer/quine.py | 43 +++++++++++++------ 30 files changed, 230 insertions(+), 121 deletions(-) create mode 100644 .github/workflows/linting-ruff.yml rename skellytracker/trackers/{image_demo_viewer => demo_viewers}/image_demo_viewer.py (82%) rename skellytracker/trackers/{webcam_demo_viewer => demo_viewers}/webcam_demo_viewer.py (97%) diff --git a/.github/workflows/linting-ruff.yml b/.github/workflows/linting-ruff.yml new file mode 100644 index 0000000..5827a66 --- /dev/null +++ b/.github/workflows/linting-ruff.yml @@ -0,0 +1,27 @@ +name: Linting Ruff + +on: + pull_request: + branches: [ main ] + paths: + - 'skellytracker/**' + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python 3.x + uses: actions/setup-python@v4 + with: + # Semantic version range syntax or exact version of a Python version + python-version: '3.9' + # Optional - x64 or x86 architecture, defaults to x64 + architecture: 'x64' + cache: 'pip' + - name: Install dependencies + run: | + pip install "-e.[dev]" + - name: Run linting + run: | + ruff check skellytracker \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 85e11e6..24a0515 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,7 @@ requires-python = ">=3.9,<3.13" dynamic = ["version", "description"] [project.optional-dependencies] -dev = ["black", "bumpver", "isort", "pip-tools", "pytest"] +dev = ["black", "bumpver", "isort", "pip-tools", "pytest", "ruff"] mediapipe = ["mediapipe==0.10.14"] yolo = ["ultralytics~=8.0.202"] all = ["ultralytics~=8.0.202", "mediapipe==0.10.14"] @@ -85,6 +85,12 @@ push = true [tool.bumpver.file_patterns] "skellytracker/__init__.py" = ["{version}"] +[tool.ruff.lint.per-file-ignores] +"*/tests/*" = ["S101"] + +[tool.ruff.lint] +extend-select = ["B", "S", "C4", "ISC", "PERF"] + [project.scripts] skellytracker = "skellytracker.__main__:cli_main" diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index e09153e..aa18700 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -5,11 +5,12 @@ BrightestPointTracker, ) from skellytracker.trackers.charuco_tracker.charuco_tracker import CharucoTracker + try: from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( MediapipeHolisticTracker, ) -except: +except ModuleNotFoundError: print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker @@ -19,7 +20,7 @@ from skellytracker.trackers.segment_anything_tracker.segment_anything_tracker import ( SAMTracker, ) -except: +except ModuleNotFoundError: print("To use yolo_tracker, install skellytracker[yolo]") @@ -29,15 +30,14 @@ def main(demo_tracker: str = "mediapipe_holistic_tracker"): BrightestPointTracker(num_points=2).demo() elif demo_tracker == "charuco_tracker": - charuco_squares_x_in = 7 - charuco_squares_y_in = 5 - number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_squares_x = 7 + charuco_squares_y = 5 + number_of_charuco_markers = (charuco_squares_x - 1) * (charuco_squares_y - 1) charuco_ids = [str(index) for index in range(number_of_charuco_markers)] - CharucoTracker( tracked_object_names=charuco_ids, - squares_x=charuco_squares_x_in, - squares_y=charuco_squares_y_in, + squares_x=charuco_squares_x, + squares_y=charuco_squares_y, dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), ).demo() @@ -56,6 +56,7 @@ def main(demo_tracker: str = "mediapipe_holistic_tracker"): SAMTracker().demo() elif demo_tracker == "yolo_object_tracker": YOLOObjectTracker(model_size="medium").demo() - + + if __name__ == "__main__": - main(demo_tracker="mediapipe_holistic_tracker") \ No newline at end of file + main(demo_tracker="mediapipe_holistic_tracker") diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index 02f112b..5ef6a2c 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -10,11 +10,11 @@ from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( MediapipeHolisticTracker, ) -except: +except ModuleNotFoundError: print("\n\nTo use mediapipe_holistic_tracker, install skellytracker[mediapipe]\n\n") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker -except: +except ModuleNotFoundError: print("\n\nTo use yolo_tracker, install skellytracker[yolo]\n\n") @@ -26,15 +26,14 @@ BrightestPointTracker(num_points=2).image_demo(image_path=image_path) elif demo_tracker == "charuco_tracker": - charuco_squares_x_in = 7 - charuco_squares_y_in = 5 - number_of_charuco_markers = (charuco_squares_x_in - 1) * (charuco_squares_y_in - 1) + charuco_squares_x = 7 + charuco_squares_y = 5 + number_of_charuco_markers = (charuco_squares_x - 1) * (charuco_squares_y - 1) charuco_ids = [str(index) for index in range(number_of_charuco_markers)] - CharucoTracker( tracked_object_names=charuco_ids, - squares_x=charuco_squares_x_in, - squares_y=charuco_squares_y_in, + squares_x=charuco_squares_x, + squares_y=charuco_squares_y, dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), ).image_demo(image_path=image_path) diff --git a/skellytracker/__init__.py b/skellytracker/__init__.py index 91413a8..b883a71 100644 --- a/skellytracker/__init__.py +++ b/skellytracker/__init__.py @@ -6,9 +6,13 @@ __author__ = """Skelly FreeMoCap""" __email__ = "info@freemocap.org" __repo_owner_github_user_name__ = "freemocap" -__repo_url__ = f"https://github.com/{__repo_owner_github_user_name__}/{__package_name__}" +__repo_url__ = ( + f"https://github.com/{__repo_owner_github_user_name__}/{__package_name__}" +) __repo_issues_url__ = f"{__repo_url__}/issues" +# ruff: noqa: F401, E402 + import sys from pathlib import Path @@ -24,20 +28,27 @@ from skellytracker.system.logging_configuration import configure_logging try: - from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import MediapipeHolisticTracker - from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import MediapipeModelInfo -except: + from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( + MediapipeHolisticTracker, + ) + from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import ( + MediapipeModelInfo, + ) +except ModuleNotFoundError: print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOModelInfo -except: +except ModuleNotFoundError: print("To use yolo_tracker, install skellytracker[yolo]") try: - from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import YOLOMediapipeComboTracker -except: - print("To use yolo_mediapipe_combo_tracker, install skellytracker[mediapipe, yolo] or skellytracker[all]") - + from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import ( + YOLOMediapipeComboTracker, + ) +except ModuleNotFoundError: + print( + "To use yolo_mediapipe_combo_tracker, install skellytracker[mediapipe, yolo] or skellytracker[all]" + ) -configure_logging(log_file_path=str(get_log_file_path())) \ No newline at end of file +configure_logging(log_file_path=str(get_log_file_path())) diff --git a/skellytracker/__main__.py b/skellytracker/__main__.py index eaf98ec..bb9f69f 100644 --- a/skellytracker/__main__.py +++ b/skellytracker/__main__.py @@ -6,20 +6,21 @@ print(f"adding base_package_path: {base_package_path} : to sys.path") sys.path.insert(0, str(base_package_path)) # add parent directory to sys.path -import logging +import logging # noqa: E402 -logger = logging.getLogger(__name__) +from skellytracker.RUN_ME import main # noqa: E402 -from skellytracker.RUN_ME import main +logger = logging.getLogger(__name__) def cli_main(): - logger.info(f"Running as a script") + logger.info("Running as a script") if len(sys.argv) > 1: demo_tracker = str(sys.argv[1]) else: demo_tracker = "mediapipe_holistic_tracker" main(demo_tracker=demo_tracker) + if __name__ == "__main__": - cli_main() \ No newline at end of file + cli_main() diff --git a/skellytracker/process_folder_of_videos.py b/skellytracker/process_folder_of_videos.py index 85e6f7f..4ca1fc8 100644 --- a/skellytracker/process_folder_of_videos.py +++ b/skellytracker/process_folder_of_videos.py @@ -1,9 +1,8 @@ import logging +import numpy as np from multiprocessing import Pool, cpu_count from pathlib import Path -import sys from typing import Optional -import numpy as np from pydantic import BaseModel @@ -12,16 +11,19 @@ BrightestPointTracker, ) from skellytracker.utilities.get_video_paths import get_video_paths + try: from skellytracker.trackers.yolo_mediapipe_combo_tracker.yolo_mediapipe_combo_tracker import ( YOLOMediapipeComboTracker, ) -except: - print("\n\nTo use yolo_mediapipe_combo_tracker, install skellytracker[yolo, mediapipe]\n\n") +except ModuleNotFoundError: + print( + "\n\nTo use yolo_mediapipe_combo_tracker, install skellytracker[yolo, mediapipe]\n\n" + ) try: from skellytracker.trackers.yolo_tracker.yolo_tracker import YOLOPoseTracker from skellytracker.trackers.yolo_tracker.yolo_model_info import YOLOTrackingParams -except: +except ModuleNotFoundError: print("To use yolo_tracker, install skellytracker[yolo]") try: from skellytracker.trackers.mediapipe_tracker.mediapipe_holistic_tracker import ( @@ -30,7 +32,7 @@ from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import ( MediapipeTrackingParams, ) -except: +except ModuleNotFoundError: print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") logger = logging.getLogger(__name__) @@ -113,7 +115,7 @@ def process_single_video( tracking_params: BaseModel, video_path: Path, annotated_video_path: Path, -) -> np.ndarray: +) -> Optional[np.ndarray]: """ Process a single video with the given tracker. Tracked data will be saved to a .npy file with the shape (numCams, numFrames, numTrackedPoints, pixelXYZ). @@ -180,11 +182,12 @@ def get_tracker(tracker_name: str, tracking_params: BaseModel) -> BaseTracker: return tracker + def get_tracker_params(tracker_name: str) -> BaseModel: if tracker_name == "MediapipeHolisticTracker": return MediapipeTrackingParams() elif tracker_name == "YOLOMediapipeComboTracker": - return YOLOTrackingParams() + return YOLOTrackingParams() # TODO: figure out how to reference both tracking params in a stable way elif tracker_name == "YOLOPoseTracker": return YOLOTrackingParams() elif tracker_name == "BrightestPointTracker": diff --git a/skellytracker/system/default_paths.py b/skellytracker/system/default_paths.py index a570c60..ca0127f 100644 --- a/skellytracker/system/default_paths.py +++ b/skellytracker/system/default_paths.py @@ -10,16 +10,24 @@ FIGSHARE_TEST_IMAGE_URL = "https://figshare.com/ndownloader/files/47043898" FIGSHARE_CHARUCO_TEST_IMAGE_URL = "https://figshare.com/ndownloader/files/47127685" + def get_base_folder_path(): - base_folder = Path().home() / BASE_FOLDER_NAME + base_folder = Path().home() / BASE_FOLDER_NAME base_folder.mkdir(exist_ok=True, parents=True) return base_folder + def get_log_file_path(): - log_file_path = get_base_folder_path() / LOGS_INFO_AND_SETTINGS_FOLDER_NAME / LOG_FILE_FOLDER_NAME / create_log_file_name() + log_file_path = ( + get_base_folder_path() + / LOGS_INFO_AND_SETTINGS_FOLDER_NAME + / LOG_FILE_FOLDER_NAME + / create_log_file_name() + ) log_file_path.parent.mkdir(exist_ok=True, parents=True) return log_file_path + def create_log_file_name(): return "log_" + get_iso6201_time_string() + ".log" @@ -29,11 +37,14 @@ def get_gmt_offset_string(): gmt_offset_int = int(time.localtime().tm_gmtoff / 60 / 60) return f"{gmt_offset_int:+}" -def get_iso6201_time_string(timespec: str = "milliseconds", make_filename_friendly: bool = True): + +def get_iso6201_time_string( + timespec: str = "milliseconds", make_filename_friendly: bool = True +): iso6201_timestamp = datetime.now().isoformat(timespec=timespec) gmt_offset_string = f"_gmt{get_gmt_offset_string()}" iso6201_timestamp_w_gmt = iso6201_timestamp + gmt_offset_string if make_filename_friendly: iso6201_timestamp_w_gmt = iso6201_timestamp_w_gmt.replace(":", "_") iso6201_timestamp_w_gmt = iso6201_timestamp_w_gmt.replace(".", "ms") - return iso6201_timestamp_w_gmt \ No newline at end of file + return iso6201_timestamp_w_gmt diff --git a/skellytracker/system/logging_configuration.py b/skellytracker/system/logging_configuration.py index a2d7bd0..b9c8b58 100644 --- a/skellytracker/system/logging_configuration.py +++ b/skellytracker/system/logging_configuration.py @@ -5,7 +5,6 @@ from typing import Optional - DEFAULT_LOGGING = {"version": 1, "disable_existing_loggers": False} @@ -40,4 +39,4 @@ def configure_logging(log_file_path: Optional[str] = ""): logger.info(f"Added logging handlers: {handlers}") else: logger = logging.getLogger(__name__) - logger.info("Logging already configured!") \ No newline at end of file + logger.info("Logging already configured!") diff --git a/skellytracker/tests/conftest.py b/skellytracker/tests/conftest.py index 3c4ca48..aa0a87e 100644 --- a/skellytracker/tests/conftest.py +++ b/skellytracker/tests/conftest.py @@ -8,14 +8,16 @@ class SessionInfo: test_image: np.ndarray charuco_test_image: np.ndarray + def pytest_sessionstart(session): SessionInfo.test_image = download_test_image() SessionInfo.charuco_test_image = download_test_image(test_image_url=FIGSHARE_CHARUCO_TEST_IMAGE_URL) -@pytest.fixture + +@pytest.fixture() def test_image(): return SessionInfo.test_image @pytest.fixture def charuco_test_image(): - return SessionInfo.charuco_test_image \ No newline at end of file + return SessionInfo.charuco_test_image diff --git a/skellytracker/tests/test_brightest_point_tracker.py b/skellytracker/tests/test_brightest_point_tracker.py index abe97e6..8652ca7 100644 --- a/skellytracker/tests/test_brightest_point_tracker.py +++ b/skellytracker/tests/test_brightest_point_tracker.py @@ -8,7 +8,7 @@ ) -@pytest.fixture +@pytest.fixture() def sample_image(): """ Create a sample image with bright spots for testing. @@ -66,11 +66,16 @@ def test_annotate_image(sample_image): bright_point_0 = tracked_objects["brightest_point_0"] bright_point_1 = tracked_objects["brightest_point_1"] + assert bright_point_0.pixel_x is not None + assert bright_point_0.pixel_y is not None + assert bright_point_1.pixel_x is not None + assert bright_point_1.pixel_y is not None + assert tracker.annotated_image[ - bright_point_0.pixel_y, bright_point_0.pixel_x + int(bright_point_0.pixel_y), int(bright_point_0.pixel_x) ].tolist() == [0, 0, 255] assert tracker.annotated_image[ - bright_point_1.pixel_y, bright_point_1.pixel_x + int(bright_point_1.pixel_y), int(bright_point_1.pixel_x) ].tolist() == [0, 0, 255] diff --git a/skellytracker/tests/test_charuco_tracker.py b/skellytracker/tests/test_charuco_tracker.py index 235b77c..3ffe0d4 100644 --- a/skellytracker/tests/test_charuco_tracker.py +++ b/skellytracker/tests/test_charuco_tracker.py @@ -1,4 +1,3 @@ -import math import cv2 import pytest import numpy as np diff --git a/skellytracker/tests/test_yolo_object_tracker.py b/skellytracker/tests/test_yolo_object_tracker.py index 4cd0843..4c20e81 100644 --- a/skellytracker/tests/test_yolo_object_tracker.py +++ b/skellytracker/tests/test_yolo_object_tracker.py @@ -15,9 +15,14 @@ def test_process_image_person_only(test_image): assert len(tracked_objects) == 1 assert tracked_objects["object"] is not None assert tracked_objects["object"].extra["boxes_xyxy"] is not None - assert np.allclose(tracked_objects["object"].extra["boxes_xyxy"], [90.676,96.981,493.54,812.03], atol=1e-2) + assert np.allclose( + tracked_objects["object"].extra["boxes_xyxy"], + [90.676, 96.981, 493.54, 812.03], + atol=1e-2, + ) assert tracked_objects["object"].extra["original_image_shape"] == (1280, 720) + @pytest.mark.usefixtures("test_image") def test_annotate_image(test_image): tracker = YOLOObjectTracker() @@ -35,8 +40,6 @@ def test_record(test_image): processed_results = tracker.recorder.process_tracked_objects() assert processed_results is not None - assert processed_results.shape == (1,4) + assert processed_results.shape == (1, 4) - assert np.allclose( - processed_results, [90.676,96.981,493.54,812.03], atol=1e-2 - ) + assert np.allclose(processed_results, [90.676, 96.981, 493.54, 812.03], atol=1e-2) diff --git a/skellytracker/tests/test_yolo_pose_tracker.py b/skellytracker/tests/test_yolo_pose_tracker.py index b0c5343..bbe31e2 100644 --- a/skellytracker/tests/test_yolo_pose_tracker.py +++ b/skellytracker/tests/test_yolo_pose_tracker.py @@ -92,7 +92,9 @@ def test_record(test_image): ] ] ) - assert np.allclose(processed_results[:, :, :2], expected_results[:, :, :2], atol=1e-2) + assert np.allclose( + processed_results[:, :, :2], expected_results[:, :, :2], atol=1e-2 + ) assert np.isnan(processed_results[:, :, 2]).all() @@ -112,6 +114,10 @@ def test_unpack_empty_results(): tracked_person = tracker.tracked_objects["tracked_person"] assert tracked_person.pixel_x is None assert tracked_person.pixel_y is None - assert tracked_person.extra["landmarks"].shape == (1, YOLOModelInfo.num_tracked_points, 2) + assert tracked_person.extra["landmarks"].shape == ( + 1, + YOLOModelInfo.num_tracked_points, + 2, + ) assert np.isnan(tracked_person.extra["landmarks"][0, :, 0]).all() assert np.isnan(tracked_person.extra["landmarks"][0, :, 1]).all() diff --git a/skellytracker/trackers/base_tracker/base_tracker.py b/skellytracker/trackers/base_tracker/base_tracker.py index a21c0ed..4fd9913 100644 --- a/skellytracker/trackers/base_tracker/base_tracker.py +++ b/skellytracker/trackers/base_tracker/base_tracker.py @@ -10,8 +10,8 @@ from skellytracker.trackers.base_tracker.base_recorder import BaseRecorder from skellytracker.trackers.base_tracker.tracked_object import TrackedObject from skellytracker.trackers.base_tracker.video_handler import VideoHandler -from skellytracker.trackers.image_demo_viewer.image_demo_viewer import ImageDemoViewer -from skellytracker.trackers.webcam_demo_viewer.webcam_demo_viewer import ( +from skellytracker.trackers.demo_viewers.image_demo_viewer import ImageDemoViewer +from skellytracker.trackers.demo_viewers.webcam_demo_viewer import ( WebcamDemoViewer, ) @@ -26,7 +26,7 @@ class BaseTracker(ABC): def __init__( self, recorder: BaseRecorder, - tracked_object_names: List[str] = [], + tracked_object_names: List[str], **data: Any, ): self.recorder = recorder @@ -112,7 +112,7 @@ def process_video( logger.error( f"Failed to load an image from: {str(input_video_filepath)}" ) - raise Exception + raise ValueError("Failed to load an image from: " + str(input_video_filepath)) self.process_image(frame) if self.recorder is not None: diff --git a/skellytracker/trackers/base_tracker/tracked_object.py b/skellytracker/trackers/base_tracker/tracked_object.py index 7daa88a..daf3a13 100644 --- a/skellytracker/trackers/base_tracker/tracked_object.py +++ b/skellytracker/trackers/base_tracker/tracked_object.py @@ -7,8 +7,9 @@ class TrackedObject: """ A dataclass for storing information about a tracked object in a single image/frame """ + object_id: str pixel_x: Optional[float] = None pixel_y: Optional[float] = None depth_z: Optional[float] = None - extra: Optional[Dict[str, Any]] = field(default_factory=dict) \ No newline at end of file + extra: Dict[str, Any] = field(default_factory=dict) diff --git a/skellytracker/trackers/base_tracker/video_handler.py b/skellytracker/trackers/base_tracker/video_handler.py index 13cc064..9a2087a 100644 --- a/skellytracker/trackers/base_tracker/video_handler.py +++ b/skellytracker/trackers/base_tracker/video_handler.py @@ -24,10 +24,8 @@ def __init__( :param codec: The codec to use for the output video. """ self.output_path = output_path - fourcc = cv2.VideoWriter_fourcc(*codec) - self.video_writer = cv2.VideoWriter( - str(output_path), fourcc, fps, frame_size - ) + fourcc = cv2.VideoWriter.fourcc(*codec) + self.video_writer = cv2.VideoWriter(str(output_path), fourcc, fps, frame_size) def add_frame(self, frame: np.ndarray) -> None: """ diff --git a/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py b/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py index 7314538..f550a3d 100644 --- a/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py +++ b/skellytracker/trackers/bright_point_tracker/brightest_point_recorder.py @@ -1,5 +1,4 @@ from typing import Dict -from copy import deepcopy import numpy as np from skellytracker.trackers.base_tracker.base_recorder import BaseRecorder diff --git a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py index a5f529d..697bf4f 100644 --- a/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py +++ b/skellytracker/trackers/bright_point_tracker/brightest_point_tracker.py @@ -71,10 +71,14 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] for i, patch in enumerate(largest_patches): self.tracked_objects[f"brightest_point_{i}"].pixel_x = patch.centroid_x self.tracked_objects[f"brightest_point_{i}"].pixel_y = patch.centroid_y - self.tracked_objects[f"brightest_point_{i}"].extra["thresholdedimage"] = thresholded_image + self.tracked_objects[f"brightest_point_{i}"].extra[ + "thresholdedimage" + ] = thresholded_image for i in range(len(largest_patches), self.num_points): - self.tracked_objects[f"brightest_point_{i}"].pixel_x = None # TODO: Is this the right value for missing data? + self.tracked_objects[f"brightest_point_{i}"].pixel_x = ( + None # TODO: Is this the right value for missing data? + ) self.tracked_objects[f"brightest_point_{i}"].pixel_y = None self.annotated_image = self.annotate_image( diff --git a/skellytracker/trackers/charuco_tracker/charuco_tracker.py b/skellytracker/trackers/charuco_tracker/charuco_tracker.py index 62d7435..3f4d50a 100644 --- a/skellytracker/trackers/charuco_tracker/charuco_tracker.py +++ b/skellytracker/trackers/charuco_tracker/charuco_tracker.py @@ -8,15 +8,15 @@ from skellytracker.trackers.charuco_tracker.charuco_recorder import CharucoRecorder +default_aruco_dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250) + class CharucoTracker(BaseTracker): def __init__( self, tracked_object_names: List[str], squares_x: int, squares_y: int, - dictionary: cv2.aruco.Dictionary = cv2.aruco.getPredefinedDictionary( - cv2.aruco.DICT_4X4_250 - ), + dictionary: cv2.aruco.Dictionary = default_aruco_dictionary, square_length: float = 1, marker_length: float = 0.8, ): diff --git a/skellytracker/trackers/image_demo_viewer/image_demo_viewer.py b/skellytracker/trackers/demo_viewers/image_demo_viewer.py similarity index 82% rename from skellytracker/trackers/image_demo_viewer/image_demo_viewer.py rename to skellytracker/trackers/demo_viewers/image_demo_viewer.py index 7fd40be..2cdd0c6 100644 --- a/skellytracker/trackers/image_demo_viewer/image_demo_viewer.py +++ b/skellytracker/trackers/demo_viewers/image_demo_viewer.py @@ -1,11 +1,13 @@ from pathlib import Path +from typing import Optional import cv2 # Constants for key actions KEY_QUIT = ord("q") + class ImageDemoViewer: - def __init__(self, tracker, window_title: str = None): + def __init__(self, tracker, window_title: Optional[str] = None): """ Initialize with a tracker and optional window title and default exposure. """ @@ -20,7 +22,7 @@ def run(self, image_path: Path): """ image = cv2.imread(str(image_path)) - tracked_results = self.tracker.process_image(image) + self.tracker.process_image(image) annotated_image = self.tracker.annotated_image diff --git a/skellytracker/trackers/webcam_demo_viewer/webcam_demo_viewer.py b/skellytracker/trackers/demo_viewers/webcam_demo_viewer.py similarity index 97% rename from skellytracker/trackers/webcam_demo_viewer/webcam_demo_viewer.py rename to skellytracker/trackers/demo_viewers/webcam_demo_viewer.py index a0daa2d..65ee82d 100644 --- a/skellytracker/trackers/webcam_demo_viewer/webcam_demo_viewer.py +++ b/skellytracker/trackers/demo_viewers/webcam_demo_viewer.py @@ -1,4 +1,5 @@ import logging +from typing import Optional import cv2 logger = logging.getLogger(__name__) @@ -18,7 +19,7 @@ def __init__( self, tracker, recorder=None, - window_title: str = None, + window_title: Optional[str] = None, default_exposure: int = DEFAULT_EXPOSURE, ): """ diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py index de8c9cb..f6254b2 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py @@ -21,7 +21,9 @@ def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: def process_tracked_objects(self, **kwargs) -> np.ndarray: image_size = kwargs.get("image_size") if image_size is None: - raise ValueError(f"image_size must be provided to process tracked objects from {__class__.__name__}") + raise ValueError( + f"image_size must be provided to process tracked objects from {__class__.__name__}" + ) self.recorded_objects_array = np.zeros( ( len(self.recorded_objects), diff --git a/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py b/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py index abcf4c1..aec4020 100644 --- a/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py +++ b/skellytracker/trackers/mmpose_tracker/mmpose_tracker.py @@ -1,33 +1,27 @@ -import cv2 -import numpy as np from mmpose.apis import inference_top_down_pose_model, init_pose_model from skellytracker.trackers.base_tracker.base_tracker import BaseTracker # correct/fill this out based on these docs: https://github.com/open-mmlab/mmpose/blob/main/docs/en/user_guides/inference.md + class MMPoseTracker(BaseTracker): def __init__(self, config_file, checkpoint_file): super().__init__(recorder=None, tracked_object_names=["human_pose"]) - self.model = init_pose_model(config_file, checkpoint_file, device='cuda:0') + self.model = init_pose_model(config_file, checkpoint_file, device="cuda:0") def process_image(self, image, **kwargs): pose_results, returned_outputs = inference_top_down_pose_model( self.model, image, bbox_thr=None, - format='xyxy', - dataset='TopDownCocoDataset', + format="xyxy", + dataset="TopDownCocoDataset", return_heatmap=False, - outputs=None + outputs=None, ) # Draw the poses on the image - self.model.show_result( - image, - pose_results, - show=False, - out_file=None - ) + self.model.show_result(image, pose_results, show=False, out_file=None) # Update the tracking data self.tracking_data = {"human_pose": pose_results} @@ -39,7 +33,8 @@ def process_image(self, image, **kwargs): "raw_image": image, } + if __name__ == "__main__": - config_file = 'configs/top_down/hrnet/coco/hrnet_w48_coco_256x192.py' - checkpoint_file = 'https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth' + config_file = "configs/top_down/hrnet/coco/hrnet_w48_coco_256x192.py" + checkpoint_file = "https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth" MMPoseTracker(config_file, checkpoint_file).demo() diff --git a/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py b/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py index f9a8882..fceffb9 100644 --- a/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py +++ b/skellytracker/trackers/segment_anything_tracker/segment_anything_tracker.py @@ -1,20 +1,21 @@ -import cv2 import numpy as np -from typing import Dict from ultralytics import SAM from skellytracker.trackers.base_tracker.base_tracker import BaseTracker + class SAMTracker(BaseTracker): def __init__(self): - super().__init__(recorder=None,tracked_object_names=["segmentation"]) + super().__init__(recorder=None, tracked_object_names=["segmentation"]) - self.model = SAM('sam_b.pt') + self.model = SAM("sam_b.pt") def process_image(self, image, **kwargs): results = self.model.predict(image) - self.tracked_objects["segmentation"].extra["landmarks"] = np.array(results[0].keypoints) + self.tracked_objects["segmentation"].extra["landmarks"] = np.array( + results[0].keypoints + ) self.annotated_image = self.annotate_image(image, results=results, **kwargs) @@ -22,7 +23,7 @@ def process_image(self, image, **kwargs): def annotate_image(self, image: np.ndarray, results, **kwargs) -> np.ndarray: return results[0].plot() - + if __name__ == "__main__": SAMTracker().demo() diff --git a/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py b/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py index 0257bbf..f6114dd 100644 --- a/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py +++ b/skellytracker/trackers/yolo_object_tracker/yolo_object_tracker.py @@ -32,9 +32,17 @@ def __init__( self.classes = None # None includes all classes def process_image(self, image, **kwargs) -> Dict[str, TrackedObject]: - results = self.model(image, classes=self.classes, max_det=1, verbose=False, conf=self.confidence_threshold) + results = self.model( + image, + classes=self.classes, + max_det=1, + verbose=False, + conf=self.confidence_threshold, + ) - box_xyxy = np.asarray(results[0].boxes.xyxy.cpu()).flatten() # On GPU, need to copy to CPU before np array conversion + box_xyxy = np.asarray( + results[0].boxes.xyxy.cpu() + ).flatten() # On GPU, need to copy to CPU before np array conversion if box_xyxy.size > 0: self.tracked_objects["object"].pixel_x = (box_xyxy[0] + box_xyxy[2]) / 0.5 diff --git a/skellytracker/trackers/yolo_tracker/yolo_tracker.py b/skellytracker/trackers/yolo_tracker/yolo_tracker.py index 8bce06b..2e7cc35 100644 --- a/skellytracker/trackers/yolo_tracker/yolo_tracker.py +++ b/skellytracker/trackers/yolo_tracker/yolo_tracker.py @@ -21,7 +21,9 @@ def process_image(self, image: np.ndarray, **kwargs) -> Dict[str, TrackedObject] self.unpack_results(results) - self.annotated_image = self.annotate_image(image=image, results=results, **kwargs) + self.annotated_image = self.annotate_image( + image=image, results=results, **kwargs + ) return self.tracked_objects @@ -35,8 +37,12 @@ def unpack_results(self, results: list): ) if tracked_person.size != 0: # add averages of all tracked points as pixel x and y - self.tracked_objects["tracked_person"].pixel_x = np.mean(tracked_person[:, 0]) - self.tracked_objects["tracked_person"].pixel_y = np.mean(tracked_person[:, 1]) + self.tracked_objects["tracked_person"].pixel_x = float( + np.mean(tracked_person[:, 0]) + ) + self.tracked_objects["tracked_person"].pixel_y = float( + np.mean(tracked_person[:, 1]) + ) self.tracked_objects["tracked_person"].extra["landmarks"] = tracked_person else: self.tracked_objects["tracked_person"].pixel_x = None diff --git a/skellytracker/utilities/download_test_image.py b/skellytracker/utilities/download_test_image.py index 4680827..3c290d4 100644 --- a/skellytracker/utilities/download_test_image.py +++ b/skellytracker/utilities/download_test_image.py @@ -1,5 +1,4 @@ import logging -from pathlib import Path import cv2 import numpy as np @@ -10,6 +9,7 @@ logger = logging.getLogger(__name__) + def download_test_image(test_image_url: str = FIGSHARE_TEST_IMAGE_URL) -> np.ndarray: try: logger.info(f"Downloading test image from {test_image_url}...") @@ -20,7 +20,7 @@ def download_test_image(test_image_url: str = FIGSHARE_TEST_IMAGE_URL) -> np.nda image_array = np.frombuffer(r.content, np.uint8) image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) - logger.info(f"Test image downloaded successfully.") + logger.info("Test image downloaded successfully.") return image except requests.exceptions.RequestException as e: @@ -29,4 +29,4 @@ def download_test_image(test_image_url: str = FIGSHARE_TEST_IMAGE_URL) -> np.nda if __name__ == "__main__": - test_data_path = download_test_image() \ No newline at end of file + test_data_path = download_test_image() diff --git a/skellytracker/utilities/get_video_paths.py b/skellytracker/utilities/get_video_paths.py index b369f8c..484cff2 100644 --- a/skellytracker/utilities/get_video_paths.py +++ b/skellytracker/utilities/get_video_paths.py @@ -18,4 +18,4 @@ def get_unique_list(list: list) -> list: unique_list = [] [unique_list.append(element) for element in list if element not in unique_list] - return unique_list \ No newline at end of file + return unique_list diff --git a/skellytracker/utilities/quine_directory_printer/quine.py b/skellytracker/utilities/quine_directory_printer/quine.py index 954cb6f..2d526af 100644 --- a/skellytracker/utilities/quine_directory_printer/quine.py +++ b/skellytracker/utilities/quine_directory_printer/quine.py @@ -34,12 +34,19 @@ class Quine: The source code is enclosed in ```python ``` code blocks. """ - def __init__(self, base_directory: str, excluded_directories: List[str], included_extensions: List[str]): + def __init__( + self, + base_directory: str, + excluded_directories: List[str], + included_extensions: List[str], + ): self.base_directory = base_directory self.excluded_directories = excluded_directories self.included_extensions = included_extensions - def write_to_file(self, root_directory: str, file_name: str, output_file: object) -> None: + def write_to_file( + self, root_directory: str, file_name: str, output_file: object + ) -> None: """ Writes the content of the file to the output markdown file. @@ -69,15 +76,24 @@ def generate_quine(self) -> None: current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") file_name = f"quine_{current_time}.md" output_dir = Path().cwd() / "output" # Output directory - output_dir.mkdir(parents=True, exist_ok=True) # Create the output directory if it doesn't exist + output_dir.mkdir( + parents=True, exist_ok=True + ) # Create the output directory if it doesn't exist file_path = output_dir / file_name # Output file path with open(file_path, "w") as output_file: for root_directory, directories, files in os.walk(self.base_directory): - directories[:] = [directory for directory in directories if directory not in self.excluded_directories] + directories[:] = [ + directory + for directory in directories + if directory not in self.excluded_directories + ] if root_directory != ".": output_file.write(f"# {os.path.relpath(root_directory, '..')}\n\n") for file_name in files: - if any(file_name.endswith(extension) for extension in self.included_extensions): + if any( + file_name.endswith(extension) + for extension in self.included_extensions + ): self.write_to_file(root_directory, file_name, output_file) @@ -86,11 +102,14 @@ def generate_quine(self) -> None: base_directory_in = r"C:\Users\jonma\github_repos\freemocap_organization\skelly_tracker\skelly_tracker\trackers" quine = Quine( base_directory=base_directory_in, - excluded_directories=["__pycache__", - ".git", - "output", - # "mediapipe_tracker", - # "charuco_tracker", - "mmpose_tracker"], - included_extensions=[".py"]) + excluded_directories=[ + "__pycache__", + ".git", + "output", + # "mediapipe_tracker", + # "charuco_tracker", + "mmpose_tracker", + ], + included_extensions=[".py"], + ) quine.generate_quine() From dead477e03b08ae30c4d0f0231d1fa9198dd7c2b Mon Sep 17 00:00:00 2001 From: Aaron Cherian Date: Wed, 28 Aug 2024 11:14:42 -0400 Subject: [PATCH 21/21] Aaron/openpose tracking (#47) Co-authored-by: philipqueen --- skellytracker/RUN_ME.py | 10 +- skellytracker/SINGLE_IMAGE_RUN.py | 14 +- skellytracker/process_folder_of_videos.py | 43 ++++- .../trackers/base_tracker/base_recorder.py | 23 ++- .../trackers/base_tracker/base_tracker.py | 55 +++++- .../openpose_tracker/openpose_model_info.py | 164 ++++++++++++++++++ .../openpose_tracker/openpose_recorder.py | 116 +++++++++++++ .../openpose_tracker/openpose_tracker.py | 149 ++++++++++++++++ 8 files changed, 558 insertions(+), 16 deletions(-) create mode 100644 skellytracker/trackers/openpose_tracker/openpose_model_info.py create mode 100644 skellytracker/trackers/openpose_tracker/openpose_recorder.py create mode 100644 skellytracker/trackers/openpose_tracker/openpose_tracker.py diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index 53e3533..30084a1 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -22,9 +22,15 @@ def main(demo_tracker: str = "mediapipe_holistic_tracker"): BrightestPointTracker().demo() elif demo_tracker == "charuco_tracker": + charuco_squares_x = 7 + charuco_squares_y = 5 + number_of_charuco_markers = (charuco_squares_x - 1) * (charuco_squares_y - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + CharucoTracker( - squaresX=7, - squaresY=5, + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x, + squares_y=charuco_squares_y, dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), ).demo() diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index a301b11..4b0060f 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -14,9 +14,17 @@ BrightestPointTracker().image_demo(image_path=image_path) elif demo_tracker == "charuco_tracker": - CharucoTracker(squaresX=7, - squaresY=5, - dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250)).image_demo(image_path=image_path) + charuco_squares_x = 7 + charuco_squares_y = 5 + number_of_charuco_markers = (charuco_squares_x - 1) * (charuco_squares_y - 1) + charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + + CharucoTracker( + tracked_object_names=charuco_ids, + squares_x=charuco_squares_x, + squares_y=charuco_squares_y, + dictionary=cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250), + ).demo() elif demo_tracker == "mediapipe_holistic_tracker": MediapipeHolisticTracker(model_complexity=2, diff --git a/skellytracker/process_folder_of_videos.py b/skellytracker/process_folder_of_videos.py index 3fce21a..74b9455 100644 --- a/skellytracker/process_folder_of_videos.py +++ b/skellytracker/process_folder_of_videos.py @@ -1,7 +1,6 @@ import logging from multiprocessing import Pool, cpu_count from pathlib import Path -import sys from typing import Optional import numpy as np from pydantic import BaseModel @@ -38,6 +37,13 @@ except: print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]") +try: + from skellytracker.trackers.openpose_tracker.openpose_tracker import ( + OpenPoseTracker, + ) +except: + print("To use openpose_tracker, install skellytracker[openpose]") + logger = logging.getLogger(__name__) @@ -122,9 +128,14 @@ def process_single_video( :param annotated_video_path: Path to save annotated video to. :return: Array of tracking data """ - video_name = ( - video_path.stem + "_mediapipe.mp4" - ) # TODO: fix it so blender output doesn't require mediapipe addendum here + + if tracker_name == "OpenPoseTracker": + video_name = video_path.stem + "_openpose.avi" + else: + video_name = ( + video_path.stem + "_mediapipe.mp4" + ) # TODO: fix it so blender output doesn't require mediapipe addendum here + tracker = get_tracker(tracker_name=tracker_name, tracking_params=tracking_params) logger.info( f"Processing video: {video_name} with tracker: {tracker.__class__.__name__}" @@ -133,7 +144,7 @@ def process_single_video( input_video_filepath=video_path, output_video_filepath=annotated_video_path / video_name, save_data_bool=False, - ) + ) # TODO: raise a custom error here if output_array is None? return output_array @@ -173,6 +184,17 @@ def get_tracker(tracker_name: str, tracking_params: BaseModel) -> BaseTracker: elif tracker_name == "BrightestPointTracker": tracker = BrightestPointTracker() + elif tracker_name == "OpenPoseTracker": + tracker = OpenPoseTracker( + openpose_root_folder_path=tracking_params.openpose_root_folder_path, + output_json_folder_path=tracking_params.output_json_path, + net_resolution=tracking_params.net_resolution, + number_people_max=tracking_params.number_people_max, + track_faces=tracking_params.track_face, + track_hands=tracking_params.track_hands, + output_resolution=tracking_params.output_resolution, + ) + else: raise ValueError("Invalid tracker type") @@ -188,19 +210,26 @@ def get_tracker_params(tracker_name: str) -> BaseModel: return YOLOTrackingParams() elif tracker_name == "BrightestPointTracker": return BaseModel() + elif tracker_name == "OpenPoseTracker": + raise ValueError( + "OpenPoseTracker requires explicitly setting the OpenPose root folder path and output json path, please provide tracking params directly" + ) else: raise ValueError("Invalid tracker type") if __name__ == "__main__": + from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import MediapipeModelInfo + synchronized_video_path = Path( - "/Users/philipqueen/freemocap_data/recording_sessions/freemocap_sample_data/synchronized_videos" + "/Your/Path/To/freemocap_data/recording_sessions/freemocap_sample_data/synchronized_videos" ) + tracker_name = "YOLOMediapipeComboTracker" num_processes = None process_folder_of_videos( - tracker_name=tracker_name, + model_info=MediapipeModelInfo(), tracking_params=get_tracker_params(tracker_name=tracker_name), synchronized_video_path=synchronized_video_path, num_processes=num_processes, diff --git a/skellytracker/trackers/base_tracker/base_recorder.py b/skellytracker/trackers/base_tracker/base_recorder.py index fdb7709..a3e85a8 100644 --- a/skellytracker/trackers/base_tracker/base_recorder.py +++ b/skellytracker/trackers/base_tracker/base_recorder.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod import logging -from typing import Dict +from pathlib import Path +from typing import Dict, Union import numpy as np @@ -44,7 +45,7 @@ def clear_recorded_objects(self): self.recorded_objects = [] self.recorded_objects_array = None - def save(self, file_path: str) -> None: + def save(self, file_path: Union[str, Path]) -> None: """ Save the recorded objects to a file. @@ -52,6 +53,22 @@ def save(self, file_path: str) -> None: :return: None """ if self.recorded_objects_array is None: - self.process_tracked_objects() + self.recorded_objects_array = self.process_tracked_objects() logger.info(f"Saving recorded objects to {file_path}") np.save(file_path, self.recorded_objects_array) + + +class BaseCumulativeRecorder(BaseRecorder): + """ + A base class for recording data from cumulative trackers. + Throws a descriptive error for methods that do not apply to recording data from this type of tracker. + Trackers implementing this will only use the process_tracked_objects method to get data in the proper format. + """ + + def __init__(self): + super().__init__() + + def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: + raise NotImplementedError( + "This tracker does not support by frame recording, please use process_tracked_objects instead" + ) diff --git a/skellytracker/trackers/base_tracker/base_tracker.py b/skellytracker/trackers/base_tracker/base_tracker.py index 8f0b5a2..7b0f0f4 100644 --- a/skellytracker/trackers/base_tracker/base_tracker.py +++ b/skellytracker/trackers/base_tracker/base_tracker.py @@ -7,7 +7,7 @@ from tqdm import tqdm -from skellytracker.trackers.base_tracker.base_recorder import BaseRecorder +from skellytracker.trackers.base_tracker.base_recorder import BaseCumulativeRecorder, BaseRecorder from skellytracker.trackers.base_tracker.tracked_object import TrackedObject from skellytracker.trackers.base_tracker.video_handler import VideoHandler from skellytracker.trackers.image_demo_viewer.image_demo_viewer import ImageDemoViewer @@ -159,3 +159,56 @@ def image_demo(self, image_path: Path) -> None: image_viewer = ImageDemoViewer(self, self.__class__.__name__) image_viewer.run(image_path=image_path) + + +class BaseCumulativeTracker(BaseTracker): + """ + A base class for tracking algorithms that run cumulatively, i.e are not able to process videos frame by frame. + Throws a descriptive error for the abstract methods of BaseTracker that do not apply to this type of tracker. + Trackers inheriting from this will need to overwrite the `process_video` method. + """ + + def __init__( + self, + tracked_object_names: List[str] = [], + recorder: Optional[BaseCumulativeRecorder] = None, + **data: Any, + ): + super().__init__( + tracked_object_names=tracked_object_names, recorder=recorder, **data + ) + + def process_image(self, **kwargs) -> None: + raise NotImplementedError( + "This tracker does not support processing individual images, please use process_video instead." + ) + + def annotate_image(self, **kwargs) -> None: + raise NotImplementedError( + "This tracker does not support processing individual images, please use process_video instead." + ) + + @abstractmethod + def process_video( + self, + input_video_filepath: Union[str, Path], + output_video_filepath: Optional[Union[str, Path]] = None, + save_data_bool: bool = False, + use_tqdm: bool = True, + **kwargs, + ) -> Union[np.ndarray, None]: + """ + Run the tracker on a video. + + :param input_video_filepath: Path to video file. + :param output_video_filepath: Path to save annotated video to, does not save video if None. + :param save_data_bool: Whether to save the data to a file. + :param use_tqdm: Whether to use tqdm to show a progress bar + :return: Array of tracked keypoint data + """ + pass + + def image_demo(self, image_path: Path) -> None: + raise NotImplementedError( + "This tracker does not support processing individual images, please use process_video instead." + ) diff --git a/skellytracker/trackers/openpose_tracker/openpose_model_info.py b/skellytracker/trackers/openpose_tracker/openpose_model_info.py new file mode 100644 index 0000000..8befd3e --- /dev/null +++ b/skellytracker/trackers/openpose_tracker/openpose_model_info.py @@ -0,0 +1,164 @@ +from skellytracker.trackers.base_tracker.base_tracking_params import BaseTrackingParams +from skellytracker.trackers.base_tracker.model_info import ModelInfo + +from typing import Optional + + +class OpenPoseModelInfo(ModelInfo): + name = "openpose" + tracker_name = "OpenPoseTracker" + body_landmark_names = [ + "nose", + "neck", + "right_shoulder", + "right_elbow", + "right_wrist", + "left_shoulder", + "left_elbow", + "left_wrist", + "hip_center", + "right_hip", + "right_knee", + "right_ankle", + "left_hip", + "left_knee", + "left_ankle", + "right_eye", + "left_eye", + "right_ear", + "left_ear", + "left_big_toe", + "left_small_toe", + "left_heel", + "right_big_toe", + "right_small_toe", + "right_heel" + ] + landmark_names = body_landmark_names + num_tracked_points_body = len(body_landmark_names) + num_tracked_points_face = 70 + num_tracked_points_left_hand = 21 + num_tracked_points_right_hand = 21 + + num_tracked_points = ( + num_tracked_points_body + + num_tracked_points_left_hand + + num_tracked_points_right_hand + + num_tracked_points_face + ) + tracked_object_names = ["pose_landmarks"] + virtual_markers_definitions = { + "head_center": { + "marker_names": ["left_ear", "right_ear"], + "marker_weights": [0.5, 0.5], + }, + "trunk_center": { + "marker_names": [ + "left_shoulder", + "right_shoulder", + "left_hip", + "right_hip", + ], + "marker_weights": [0.25, 0.25, 0.25, 0.25], + }, + } + segment_connections = { + "head": {"proximal": "left_ear", "distal": "right_ear"}, + "neck": {"proximal": "head_center", "distal": "neck"}, + "spine": {"proximal": "neck", "distal": "hip_center"}, + "right_shoulder": {"proximal": "neck", "distal": "right_shoulder"}, + "left_shoulder": {"proximal": "neck", "distal": "left_shoulder"}, + "right_upper_arm": {"proximal": "right_shoulder", "distal": "right_elbow"}, + "left_upper_arm": {"proximal": "left_shoulder", "distal": "left_elbow"}, + "right_forearm": {"proximal": "right_elbow", "distal": "right_wrist"}, + "left_forearm": {"proximal": "left_elbow", "distal": "left_wrist"}, + "right_pelvis": {"proximal": "hip_center", "distal": "right_hip"}, + "left_pelvis": {"proximal": "hip_center", "distal": "left_hip"}, + "right_thigh": {"proximal": "right_hip", "distal": "right_knee"}, + "left_thigh": {"proximal": "left_hip", "distal": "left_knee"}, + "right_shank": {"proximal": "right_knee", "distal": "right_ankle"}, + "left_shank": {"proximal": "left_knee", "distal": "left_ankle"}, + "right_foot": {"proximal": "right_ankle", "distal": "right_big_toe"}, + "left_foot": {"proximal": "left_ankle", "distal": "left_big_toe"}, + "right_heel": {"proximal": "right_ankle", "distal": "right_heel"}, + "left_heel": {"proximal": "left_ankle", "distal": "left_heel"}, + "right_foot_bottom": {"proximal": "right_heel", "distal": "right_big_toe"}, + "left_foot_bottom": {"proximal": "left_heel", "distal": "left_big_toe"}, + } + center_of_mass_definitions = { #NOTE: using forearm/hand definition from Winter tables, as we don't have hand definitions here + "head": { + "segment_com_length": .5, + "segment_com_percentage": .081, + }, + "spine": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.497, + }, + "right_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "left_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "right_forearm": { + "segment_com_length": 0.682, + "segment_com_percentage": 0.022, + }, + "left_forearm": { + "segment_com_length": 0.682, + "segment_com_percentage": 0.022, + }, + "right_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "left_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "right_shank": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.0465, + }, + "left_shank": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.0465, + }, + "right_foot": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.0145, + }, + "left_foot": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.0145, + }, + } + joint_hierarchy = { + "hip_center": ["left_hip", "right_hip", "trunk_center"], + "trunk_center": ["neck"], + "neck": ["left_shoulder", "right_shoulder", "head_center"], + "head_center": ["nose", "left_ear", "right_ear", "left_eye", "right_eye"], + "left_shoulder": ["left_elbow"], + "left_elbow": ["left_wrist"], + "right_shoulder": ["right_elbow"], + "right_elbow": ["right_wrist"], + "left_hip": ["left_knee"], + "left_knee": ["left_ankle"], + "left_ankle": ["left_big_toe", "left_small_toe", "left_heel"], + "right_hip": ["right_knee"], + "right_knee": ["right_ankle"], + "right_ankle": ["right_big_toe", "right_small_toe", "right_heel"], + } + + +class OpenPoseTrackingParams(BaseTrackingParams): + openpose_root_folder_path: str + output_json_path: Optional[str] = None + net_resolution: str = "-1x320" + number_people_max: int = 1 + track_hands: bool = True + track_face: bool = True + write_video: bool = True + output_resolution: str = "-1x-1" diff --git a/skellytracker/trackers/openpose_tracker/openpose_recorder.py b/skellytracker/trackers/openpose_tracker/openpose_recorder.py new file mode 100644 index 0000000..4b57d95 --- /dev/null +++ b/skellytracker/trackers/openpose_tracker/openpose_recorder.py @@ -0,0 +1,116 @@ +import json +from typing import Dict, Union +import numpy as np +from pathlib import Path +import re +from tqdm import tqdm + +from skellytracker.trackers.base_tracker.base_recorder import BaseCumulativeRecorder +from skellytracker.trackers.openpose_tracker.openpose_model_info import ( + OpenPoseModelInfo, +) + + +class OpenPoseRecorder(BaseCumulativeRecorder): + def __init__( + self, + track_hands: bool = False, + track_faces: bool = False, + ): + super().__init__() + self.track_hands = track_hands + self.track_faces = track_faces + + def extract_frame_index(self, filename: str) -> Union[int, None]: + """Extract the numeric part indicating the frame index from the filename.""" + match = re.search(r"_(\d{12})_keypoints", filename) + return int(match.group(1)) if match else None + + def parse_openpose_jsons(self, json_directory: Union[Path, str]) -> np.ndarray: + # Remove the iteration over subdirectories and focus on a single directory + json_directory = Path(json_directory) + files = list(Path(json_directory).glob("*.json")) + num_frames = len(files) + frame_indices = [ + index + for f in files + if (index := self.extract_frame_index(f.name)) is not None + ] + frame_indices.sort() + + if len(frame_indices) != num_frames: + raise ValueError( + f"Invalid number of frames in {json_directory}: expected {num_frames} != {len(frame_indices)} frames in file" + ) + + num_markers = OpenPoseModelInfo.num_tracked_points_body + if self.track_hands: + num_markers += ( + OpenPoseModelInfo.num_tracked_points_right_hand + + OpenPoseModelInfo.num_tracked_points_left_hand + ) + if self.track_faces: + num_markers += OpenPoseModelInfo.num_tracked_points_face + + # Initialize a single camera array since we're only processing one video at a time + data_array = np.full((num_frames, num_markers, 3), np.nan) + + # Process each JSON file in the directory + for file_index, json_file in enumerate( + tqdm(files, desc=f"Processing {json_directory.name} JSONs") + ): + with open(json_file) as f: + data = json.load(f) + + if data["people"]: + keypoints = self.extract_keypoints(data["people"][0]) + data_array[frame_indices[file_index], :, :] = keypoints + + return data_array + + def extract_keypoints(self, person_data: Dict[str, np.ndarray]) -> np.ndarray: + """Extract and organize keypoints from person data.""" + + body_markers = OpenPoseModelInfo.num_tracked_points_body + hand_markers = ( + OpenPoseModelInfo.num_tracked_points_left_hand + + OpenPoseModelInfo.num_tracked_points_right_hand + ) + face_markers = OpenPoseModelInfo.num_tracked_points_face + + # Initialize a full array of NaNs for keypoints + keypoints_array = np.full( + (body_markers + (hand_markers) + face_markers, 3), np.nan + ) + + # Populate the array with available data + if "pose_keypoints_2d" in person_data: + keypoints_array[:body_markers, :] = np.reshape( + person_data["pose_keypoints_2d"], (-1, 3) + )[:body_markers, :] + if ( + "hand_left_keypoints_2d" in person_data + and "hand_right_keypoints_2d" in person_data + ): + keypoints_array[body_markers : body_markers + OpenPoseModelInfo.num_tracked_points_left_hand, :] = np.reshape( + person_data["hand_left_keypoints_2d"], (-1, 3) + )[:hand_markers, :] + keypoints_array[ + body_markers + OpenPoseModelInfo.num_tracked_points_left_hand : body_markers + OpenPoseModelInfo.num_tracked_points_left_hand + OpenPoseModelInfo.num_tracked_points_right_hand, : + ] = np.reshape(person_data["hand_right_keypoints_2d"], (-1, 3))[ + :hand_markers, : + ] + if "face_keypoints_2d" in person_data: + keypoints_array[body_markers + hand_markers :, :] = np.reshape( + person_data["face_keypoints_2d"], (-1, 3) + )[:face_markers, :] + + return keypoints_array + + def process_tracked_objects(self, output_json_path: Path) -> np.ndarray: + """ + Convert the recorded JSON data into the structured numpy array format. + """ + # In this case, the recorded_objects are already in the desired format, so we simply return them. + self.recorded_objects_array = self.parse_openpose_jsons(output_json_path) + return self.recorded_objects_array diff --git a/skellytracker/trackers/openpose_tracker/openpose_tracker.py b/skellytracker/trackers/openpose_tracker/openpose_tracker.py new file mode 100644 index 0000000..c9bd2e8 --- /dev/null +++ b/skellytracker/trackers/openpose_tracker/openpose_tracker.py @@ -0,0 +1,149 @@ +import subprocess +from pathlib import Path +from typing import Optional, Union +from skellytracker.trackers.base_tracker.base_tracker import BaseCumulativeTracker +from skellytracker.trackers.openpose_tracker.openpose_recorder import OpenPoseRecorder + + +class OpenPoseTracker(BaseCumulativeTracker): + def __init__( + self, + openpose_root_folder_path: Union[str, Path], + output_json_folder_path: Optional[Union[str, Path]] = None, + net_resolution: str = "-1x320", + number_people_max: int = 1, + track_hands: bool = True, + track_faces: bool = True, + output_resolution: str = "-1x-1", + ): + """ + Initialize the OpenPoseTracker. + + :param recorder: An instance of OpenPoseRecorder for handling the output. + :param openpose_root_folder_path: Path to the OpenPose root folder. + :param output_json_folder_path: Path to the output JSON folder. + :param net_resolution: Network resolution for OpenPose processing. + :param number_people_max: Maximum number of people to detect. + :param track_hands: Whether to track hands. + :param track_faces: Whether to track faces. + :param output_resolution: Output resolution for video. + """ + super().__init__( + tracked_object_names=[], + recorder=OpenPoseRecorder( + track_hands=track_hands, + track_faces=track_faces, + ), + track_hands=track_hands, + track_faces=track_faces, + ) + self.openpose_root_folder_path = Path(openpose_root_folder_path) + self.output_json_folder_path = output_json_folder_path + self.net_resolution = net_resolution + self.number_people_max = number_people_max + self.track_hands = track_hands + self.track_faces = track_faces + self.output_resolution = output_resolution + + def set_track_hands(self, track_hands: bool): + self._track_hands = track_hands + self.recorder.track_hands = track_hands + + def set_track_faces(self, track_faces: bool): + self._track_faces = track_faces + self.recorder.track_faces = track_faces + + def set_json_output_path(self, output_json_folder_path: Union[str, Path]): + self.output_json_folder_path = Path(output_json_folder_path) + + def process_video( + self, + input_video_filepath: Union[str, Path], + output_video_filepath: Union[str, Path], + save_data_bool: bool = False, + use_tqdm: bool = True, # TODO: this is unused, replace with an openpose flag or remove + **kwargs, + ): + """ + Run the OpenPose demo on a video file to generate JSON outputs + in a unique directory for each video. + + :param input_video_filepath: Path to the input video file. + :param output_video_filepath: Path to the output video file. + :param save_data_bool: Whether to save the data. + :param use_tqdm: Whether to use tqdm progress bar. + :return: The output array, or None if recorder isn't initialized in tracker. + """ + # Extract video name without extension to use as a unique folder name + video_name = Path(input_video_filepath).stem + + if self.output_json_folder_path is None: + self.output_json_folder_path = Path(input_video_filepath).parent.parent / "output_data" / "raw_data" / "openpose_jsons" + + Path(self.output_json_folder_path).mkdir(parents=True, exist_ok=True) + + unique_json_output_path = Path(self.output_json_folder_path) / video_name + unique_json_output_path.mkdir(parents=True, exist_ok=True) + + # Full path to the OpenPose executable + openpose_executable_path = ( + self.openpose_root_folder_path / "bin" / "OpenPoseDemo.exe" + ) + + openpose_command = [ + str(openpose_executable_path), # Full path to the OpenPose executable + "--video", + str(input_video_filepath), + "--write_json", + str(unique_json_output_path), + "--net_resolution", + str(self.net_resolution), + "--number_people_max", + str(self.number_people_max), + "--write_video", + str(output_video_filepath), + "--output_resolution", + str(self.output_resolution), + ] + + if self.track_hands: + openpose_command.append("--hand") + if self.track_faces: + openpose_command.append("--face") + + # Update the subprocess command to use the unique output directory + try: + subprocess.run( + openpose_command, + shell=False, + cwd=self.openpose_root_folder_path, # Set the current working directory for the subprocess + check=True, + ) + except subprocess.CalledProcessError as e: + print(f"Error: {e}") + return None + + if self.recorder is not None: + output_array = self.recorder.process_tracked_objects( + output_json_path=unique_json_output_path + ) + if save_data_bool: + self.recorder.save( + file_path=str(Path(input_video_filepath).with_suffix(".npy")) + ) + else: + output_array = None + + return output_array + + +if __name__ == "__main__": + # Example usage + openpose_root_folder_path = r"C:\openpose" + input_video_filepath = r'C:\path\to\input\video.mp4' + output_video_filepath = r'C:\path\to\output\video.avi' + + tracker = OpenPoseTracker( + openpose_root_folder_path=str(openpose_root_folder_path), + ) + tracker.process_video(input_video_filepath, output_video_filepath)