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/types/Skeleton.py

@@ -3,6 +3,7 @@ from enum import Enum
 import typing
 
 import cv2
+import matplotlib.pyplot as plt
 
 from .Type import Type
 from .Point3d import Point3d
@@ -12,22 +13,27 @@ 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
 class Joint(Type):
     pos: Point3d
@@ -48,14 +54,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()}