From 5b0b0f5fef95cf0d66471464607a6f08fb725943 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Wed, 24 Apr 2024 10:15:11 -0600 Subject: [PATCH 1/3] adapt yolo model info to new format --- .../trackers/yolo_tracker/yolo_model_info.py | 192 ++++++++++++------ 1 file changed, 126 insertions(+), 66 deletions(-) diff --git a/skellytracker/trackers/yolo_tracker/yolo_model_info.py b/skellytracker/trackers/yolo_tracker/yolo_model_info.py index ad7c961..dc37432 100644 --- a/skellytracker/trackers/yolo_tracker/yolo_model_info.py +++ b/skellytracker/trackers/yolo_tracker/yolo_model_info.py @@ -12,73 +12,133 @@ class YOLOModelInfo(ModelInfo): "extra_large": "yolov8x-pose.pt", "high_res": "yolov8x-pose-p6.pt", } - marker_dict = { - 0: "nose", - 1: "left_eye", - 2: "right_eye", - 3: "left_ear", - 4: "right_ear", - 5: "left_shoulder", - 6: "right_shoulder", - 7: "left_elbow", - 8: "right_elbow", - 9: "left_wrist", - 10: "right_wrist", - 11: "left_hip", - 12: "right_hip", - 13: "left_knee", - 14: "right_knee", - 15: "left_ankle", - 16: "right_ankle", - } - body_segment_names = [ - "head", - "trunk", - "right_upper_arm", - "left_upper_arm", - "right_forearm", - "left_forearm", - "right_thigh", - "left_thigh", - "right_shin", - "left_shin", - ] - joint_connections = [ - ["left_ear", "right_ear"], - ["mid_chest_marker", "mid_hip_marker"], - ["right_shoulder", "right_elbow"], - ["left_shoulder", "left_elbow"], - ["right_elbow", "right_wrist"], - ["left_elbow", "left_wrist"], - ["right_hip", "right_knee"], - ["left_hip", "left_knee"], - ["right_knee", "right_ankle"], - ["left_knee", "left_ankle"], - ] - segment_COM_lengths = [ - 0.5, - 0.5, - 0.436, - 0.436, - 0.682, - 0.682, - 0.433, - 0.433, - 0.606, - 0.606, - ] - segment_COM_percentages = [ - 0.081, - 0.497, - 0.028, - 0.028, - 0.022, - 0.022, - 0.1, - 0.1, - 0.061, - 0.061, + landmark_names = [ + "nose", + "left_eye", + "right_eye", + "left_ear", + "right_ear", + "left_shoulder", + "right_shoulder", + "left_elbow", + "right_elbow", + "left_wrist", + "right_wrist", + "left_hip", + "right_hip", + "left_knee", + "right_knee", + "left_ankle", + "right_ankle", ] + virtual_markers_definitions = { + "head_center": { + "marker_names": ["left_ear", "right_ear"], + "marker_weights": [0.5, 0.5], + }, + "neck_center": { + "marker_names": ["left_shoulder", "right_shoulder"], + "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], + }, + "hips_center": { + "marker_names": ["left_hip", "right_hip"], + "marker_weights": [0.5, 0.5], + }, + } + segment_connections = { + "head": {"proximal": "left_ear", "distal": "right_ear"}, + "neck": { + "proximal": "head_center", + "distal": "neck_center", + }, + "spine": { + "proximal": "neck_center", + "distal": "hips_center", + }, + "right_shoulder": {"proximal": "neck_center", "distal": "right_shoulder"}, + "left_shoulder": {"proximal": "neck_center", "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": "hips_center", "distal": "right_hip"}, + "left_pelvis": {"proximal": "hips_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"}, + } + center_of_mass_definitions = { + "head": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.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.606, + "segment_com_percentage": 0.061, + }, + "left_shank": { + "segment_com_length": 0.606, + "segment_com_percentage": 0.061, + }, + } + joint_hierarchy = { + "hips_center": ["left_hip", "right_hip", "trunk_center"], + "trunk_center": ["neck_center"], + "neck_center": ["left_shoulder", "right_shoulder", "head_center"], + "head_center": [ + "nose", + "left_eye", + "right_eye", + "left_ear", + "right_ear", + ], + "left_shoulder": ["left_elbow"], + "left_elbow": ["left_wrist"], + "right_shoulder": ["right_elbow"], + "right_elbow": ["right_wrist"], + "left_hip": ["left_knee"], + "left_knee": ["left_ankle"], + "right_hip": ["right_knee"], + "right_knee": ["right_ankle"], + } class YOLOTrackingParams(BaseTrackingParams): From 58187a261657e5c76407f097314863adb1770279 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Wed, 24 Apr 2024 10:15:29 -0600 Subject: [PATCH 2/3] make tracked object names optional in ModelInfo --- skellytracker/trackers/base_tracker/model_info.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skellytracker/trackers/base_tracker/model_info.py b/skellytracker/trackers/base_tracker/model_info.py index b0675ff..fb4d450 100644 --- a/skellytracker/trackers/base_tracker/model_info.py +++ b/skellytracker/trackers/base_tracker/model_info.py @@ -4,7 +4,7 @@ class ModelInfo(dict): landmark_names: List[str] num_tracked_points: int - tracked_object_names: list + tracked_object_names: Optional[list] = None virtual_markers_definitions: Optional[Dict[str, Dict[str, List[str | float]]]] = None segment_connections: Optional[Dict[str, Dict[str, str]]] = None center_of_mass_definitions: Optional[Dict[str, Dict[str, float]]] = None From 5d215612e395b74e4529669cf3fe6ad87ee6d805 Mon Sep 17 00:00:00 2001 From: philipqueen Date: Thu, 23 May 2024 12:28:21 -0600 Subject: [PATCH 3/3] add model name to model info --- skellytracker/trackers/base_tracker/model_info.py | 1 + skellytracker/trackers/yolo_tracker/yolo_model_info.py | 1 + 2 files changed, 2 insertions(+) diff --git a/skellytracker/trackers/base_tracker/model_info.py b/skellytracker/trackers/base_tracker/model_info.py index fb4d450..2ae19f7 100644 --- a/skellytracker/trackers/base_tracker/model_info.py +++ b/skellytracker/trackers/base_tracker/model_info.py @@ -2,6 +2,7 @@ class ModelInfo(dict): + model_name: str landmark_names: List[str] num_tracked_points: int tracked_object_names: Optional[list] = None diff --git a/skellytracker/trackers/yolo_tracker/yolo_model_info.py b/skellytracker/trackers/yolo_tracker/yolo_model_info.py index dc37432..ad42ac4 100644 --- a/skellytracker/trackers/yolo_tracker/yolo_model_info.py +++ b/skellytracker/trackers/yolo_tracker/yolo_model_info.py @@ -3,6 +3,7 @@ class YOLOModelInfo(ModelInfo): + model_name = "yolo" num_tracked_points = 17 model_dictionary = { "nano": "yolov8n-pose.pt",