cleanup project

dataset.py

@@ -2,6 +2,7 @@ import json
 import torch
 import os
 import numpy as np
+from utils import get_mapping_arr, apply_mapping, openpose_to_opengl_coords
 
 
 class SMPLyDataset(torch.utils.data.Dataset):
@@ -33,49 +34,14 @@ class SMPLyDataset(torch.utils.data.Dataset):
         """ 
             transform: transforms the order of an origin array to the target format
         """
-        # TODO: expand features as needed
-        # based on mappings found here
-        # https://github.com/ortegatron/playing_smplifyx/blob/master/smplifyx/utils.py
-        map = np.array([
-            # 8,  # hip - middle
-            9,  # hip - left
-            12,  # hip - right
-            -1,  # body center (belly, not present in body_25)
-            13,  # left knee
-            10,  # right knee,
-            1,  # chest
-            14,  # left ankle
-            11,  # right ankle
-            1,  # chest again ? check this one out
-            19,  # left toe
-            22,  # right toe
-            -1,  # neck (not present in body_25)
-            -1,  # between torso and left shoulder
-            -1,  # between torso and right shoulder
-            0,  # head
-            5,  # left shoulder
-            2,  # right shoulder
-            6,  # left elbow
-            3,  # right elbow
-            7,  # left hand
-            4,  # right hand
-            -1,  # left fingers
-            -1  # right fingers
 
-        ],
-            dtype=np.int32)
-        in_len = int(len(data) / 3)
-        in_data = np.array(data).reshape((in_len, 3))
-        out = np.zeros((len(map), 3), dtype=np.float32)
+        data = np.array(data).reshape((-1, 3))
 
-        for i in range(len(map)):
-            m = map[i]
-            if m == -1:
-                continue
-            # TODO: cleanup transform
-            out[i][0] = (in_data[m][0] / 1920 * 2 - 1)
-            out[i][1] = 1 - (in_data[m][1] / 1080 * 2)
-        return out
+        mapping = get_mapping_arr(origin_format, target_format)
+        # remap data to match expacted target format
+        remapped_data = apply_mapping(data, mapping)
+        # TODO: pass image resolution here
+        return openpose_to_opengl_coords(remapped_data, 1920, 1080)
 
     def __len__(self):
         # TODO: something like this could work for now we simply use one item

example.py

@@ -39,8 +39,8 @@ def main():
     r = SMPLyRenderer()
     m = l.create_model()
 
-    keypoints = dataset[0]
-    print(keypoints)
+    keypoints, conf = dataset[0]
+    print(keypoints, conf)
     betas = torch.randn([1, 10], dtype=torch.float32)
 
     # body_pose = torch.randn((1, 69), dtype=torch.float32) * 0.7

utils.py

@@ -0,0 +1,74 @@
+from typing import List, Set, Dict, Tuple, Optional
+import numpy as np
+
+
+def get_mapping_arr(
+    input_format: str = "body_25",
+    output_format: str = "smpl",
+) -> list:
+    # TODO: expand features as needed
+    # based on mappings found here
+    # https://github.com/ortegatron/playing_smplifyx/blob/master/smplifyx/utils.py
+    return np.array([
+        # 8,  # hip - middle
+        9,  # hip - left
+        12,  # hip - right
+        -1,  # body center (belly, not present in body_25)
+        13,  # left knee
+        10,  # right knee,
+        1,  # chest
+        14,  # left ankle
+        11,  # right ankle
+        1,  # chest again ? check this one out
+        19,  # left toe
+        22,  # right toe
+            -1,  # neck (not present in body_25)
+            -1,  # between torso and left shoulder
+            -1,  # between torso and right shoulder
+        0,  # head
+        5,  # left shoulder
+        2,  # right shoulder
+        6,  # left elbow
+        3,  # right elbow
+        7,  # left hand
+        4,  # right hand
+            -1,  # left fingers
+            -1  # right fingers
+    ])
+
+
+def apply_mapping(
+        input_data: List,
+        mapping: list):
+    return [input_data[i] if i != -1 else (0, 0, 0) for i in mapping]
+
+
+def openpose_to_opengl_coords(
+    input_data: List[Tuple[float, float]],
+    real_width: int,
+    real_height: int
+) -> (List[Tuple[float, float, float]], List[float]):
+    """converts a list of OpenPose 2d keypoints with confidence to a opengl coordinate system 3d point list and a confidence array
+
+    Args:
+        input_data (List[Tuple[float, float]]): [description]
+        real_width (int): OpenPose input image/data width
+        real_height (int): OpenPose input image/data height
+
+
+    Returns:
+        [type]: [description]
+    """
+
+    points = np.array([
+        [
+            x / real_width * 2 - 1,
+            y / real_height * 2,
+            0
+        ] for (x, y, z) in input_data])
+
+    conf = np.array([
+        z for (_, _, z) in input_data
+    ])
+
+    return (points, conf)