Skip to content

Commit

Permalink
Merge pull request #33 from freemocap/philip/yolo_model_info
Browse files Browse the repository at this point in the history 
Update YOLO model info
  • Loading branch information
@philipqueen
philipqueen authored May 23, 2024
2 parents e0a6c95 + 5d21561 commit 0039eae
Show file tree
Hide file tree
Showing 2 changed files with 128 additions and 67 deletions.
2 changes: 1 addition & 1 deletion skellytracker/trackers/base_tracker/model_info.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ class ModelInfo(dict):
model_name: str
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
Expand Down
193 changes: 127 additions & 66 deletions skellytracker/trackers/yolo_tracker/yolo_model_info.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@


class YOLOModelInfo(ModelInfo):
model_name = "yolo"
num_tracked_points = 17
model_dictionary = {
"nano": "yolov8n-pose.pt",
Expand All @@ -12,73 +13,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):
Expand Down

0 comments on commit 0039eae

Please sign in to comment.