Chest default

ovtk_track/output/debug.py

@@ -1,4 +1,5 @@
 import cv2
+import matplotlib.pyplot as plt
 
 from . import OutputProcess
 from ovtk_track import types
@@ -14,7 +15,10 @@ class Process(OutputProcess):
         super().__init__(*args)
 
     def setup(self):
-        pass
+        self.fig = plt.figure()
+        self.axes = self.fig.add_subplot(projection='3d')
+        self.axes.view_init(15, 0, vertical_axis='y')
+        plt.show(block=False)
 
     def send(self):
         landmarks = self._inputs['landmarks'].get_nowait()
@@ -31,8 +35,15 @@ class Process(OutputProcess):
             landmarks.draw(image, frame, label=False, color=(130, 130, 130))
 
         if skeleton is not None:
-            skeleton.draw(image, frame)
+            skeleton.draw(self.axes)
 
         cv2.imshow("face", frame)
+        plt.draw()
+
+        # event loops
+        plt.pause(0.0001)
         if cv2.waitKey(1) & 0xFF == ord('q'):
             raise KeyboardInterrupt('User requested stop')
+
+        for artist in plt.gca().lines + plt.gca().collections:
+            artist.remove()

ovtk_track/transform/solve/skel_from_3d_landmarks.py

@@ -74,6 +74,17 @@ class Process(TransformProcess):
 
                 joints[JOINT_TYPES.HEAD] = head_joint
 
+            # Synthizise other joints from existing data
+            if not joints.get(JOINT_TYPES.CHEST) and joints.get(JOINT_TYPES.HEAD):
+                chest_center = joints[JOINT_TYPES.HEAD].pos.as_np()
+                chest_center = np.power(chest_center, 3) / (1e3 + np.power(chest_center, 2))
+                chest_center -= [0, 100, 0]
+
+                chest_rot = Quaternion.identity().slerp(joints[JOINT_TYPES.HEAD].rot, 0.1)
+
+                chest_joint = Joint(Point3d(*chest_center), chest_rot)
+                joints[JOINT_TYPES.CHEST] = chest_joint
+
             skeleton = Skeleton(joints)
 
         self._outputs['skel'].send(skeleton)

ovtk_track/types/Quaternion.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 
 import numpy as np
-from scipy.spatial.transform import Rotation
+from scipy.spatial.transform import Rotation, Slerp
 from .Type import Type
 
 
@@ -12,6 +12,10 @@ class Quaternion(Type):
     y: float
     z: float
 
+    @classmethod
+    def identity(cls):
+        return cls(1, 0, 0, 0)
+
     def __mul__(self, q):
         if isinstance(q, self.__class__):
             product = self.as_np() * q.as_np()
@@ -45,6 +49,16 @@ class Quaternion(Type):
     def conjugate(self):
         return self.__class__(self.w, -self.x, -self.y, -self.z)
 
+    def slerp(self, other, t):
+        r = Rotation.from_quat([
+            [self.x, self.y, self.z, self.w],
+            [other.x, other.y, other.z, other.w],
+        ])
+
+        slerp = Slerp([0, 1], r)
+        x, y, z, w = slerp([t]).as_quat()[0]
+        return self.__class__(w, x, y, z)
+
     def draw(self, canvas, origin):
         raise NotImplementedError()
 

ovtk_track/types/Skeleton.py

@@ -2,8 +2,6 @@ from dataclasses import dataclass, field
 from enum import Enum
 import typing
 
-import cv2
-
 from .Type import Type
 from .Point3d import Point3d
 from .Quaternion import Quaternion
@@ -12,20 +10,26 @@ from .Quaternion import Quaternion
 class JOINT_TYPES(Enum):
     HEAD = 'head'
     CHEST = 'chest'
+    HIPS = 'hips'
 
     SHOULDER_L = 'shoulder_l'
     ELBOW_L = 'elbow_l'
-    WRIST_L = 'wrist_l'
-    HIP_L = 'hip_l'
     KNEE_L = 'knee_l'
     FOOT_L = 'foot_l'
+    WRIST_L = 'wrist_l'
 
     SHOULDER_R = 'shoulder_r'
     ELBOW_R = 'elbow_r'
-    WRIST_R = 'wrist_r'
-    HIP_R = 'hip_r'
     KNEE_R = 'knee_r'
     FOOT_R = 'foot_r'
+    WRIST_R = 'wrist_r'
+
+
+default_colors = [
+    [255, 0, 0], [0, 255, 0], [0, 0, 255],
+    [255, 255, 0], [255, 0, 255], [0, 255, 255],
+    [128, 255, 0], [255, 0, 128], [0, 255, 128],
+]
 
 
 @dataclass
@@ -48,14 +52,10 @@ class Skeleton(Type):
         # TODO: More intelegent merge
         return Skeleton(self.joints + other.joints)
 
-    def draw(self, image, canvas, color=(255, 255, 255)):
-        for i, joint in enumerate(self.joints.values()):
-            x, y, z = joint.pos.project_to_image(image)
-
-            if x > image.width or x < 0 or y > image.height or y < 0:
-                continue
-
-            cv2.circle(canvas, (x, y), 1, color, -1, cv2.LINE_AA)
+    def draw(self, axes, colors=default_colors):
+        xs, ys, zs = zip(*(joint.pos.as_np() for joint in self.joints.values()))
+        color = [[v / 255 for v in color] for color in colors[:len(xs)]]
+        axes.scatter(xs, ys, zs, c=color)
 
     def serialize(self):
         return {type.value: joint for type, joint in self.joints.items()}