Fix pose tracking routine

Author: vijay4313

pose_tracking/cubemos_api.py

@@ -1,5 +1,12 @@
-#!/usr/bin/env python3
-
+#!/usr/bin/env python
+# Title           :loader.py
+# Author          :Venkatraman Narayanan, Bala Murali Manoghar, Vishnu Shashank Dorbala, Aniket Bera, Dinesh Manocha
+# Copyright       :"Copyright 2020, Proxemo project"
+# Version         :1.0
+# License         :"MIT"
+# Maintainer      :Venkatraman Narayanan, Bala Murali Manoghar
+# Email           :[email protected], [email protected]
+# ==============================================================================
 import argparse
 import os
 import platform
@@ -28,6 +35,14 @@
 
 
 def default_log_dir():
+    """Get default cubemos log dir.
+
+    Raises:
+        Exception: Works only on Windows/Linux
+
+    Returns:
+        [str]: Default log directory
+    """
     if platform.system() == "Windows":
         return os.path.join(os.environ["LOCALAPPDATA"], "Cubemos", "SkeletonTracking", "logs")
     elif platform.system() == "Linux":
@@ -37,6 +52,14 @@ def default_log_dir():
 
 
 def default_license_dir():
+    """Get license file location.
+
+    Raises:
+        Exception: Works only on Windows/Linux
+
+    Returns:
+        [str]: License file directory
+    """
     if platform.system() == "Windows":
         return os.path.join(os.environ["LOCALAPPDATA"], "Cubemos", "SkeletonTracking", "license")
     elif platform.system() == "Linux":
@@ -46,6 +69,7 @@ def default_license_dir():
 
 
 def check_license_and_variables_exist():
+    """Check license file."""
     license_path = os.path.join(default_license_dir(), "cubemos_license.json")
     if not os.path.isfile(license_path):
         raise Exception(
@@ -57,13 +81,24 @@ def check_license_and_variables_exist():
         raise Exception(
             "The environment Variable \"CUBEMOS_SKEL_SDK\" is not set. "
             "Please check the troubleshooting section in the Getting "
-            "Started Guide to resolve this issue." 
+            "Started Guide to resolve this issue."
         )
 
 
 def get_valid_limbs(keypoint_ids, skeleton, confidence_threshold):
+    """Extract limbs based on confidence
+
+    Args:
+        keypoint_ids (list(tuple)): Joint ID pair (limbs)
+        skeleton (obj): Skeleton joint points
+        confidence_threshold (float): Threshold
+
+    Returns:
+        [lsit]: valid limbs (connections)
+    """
     limbs = [
-        (tuple(map(int, skeleton.joints[i])), tuple(map(int, skeleton.joints[v])))
+        (tuple(map(int, skeleton.joints[i])),
+         tuple(map(int, skeleton.joints[v])))
         for (i, v) in keypoint_ids
         if skeleton.confidences[i] >= confidence_threshold
         and skeleton.confidences[v] >= confidence_threshold
@@ -77,34 +112,17 @@ def get_valid_limbs(keypoint_ids, skeleton, confidence_threshold):
 
 
 def render_result(skeletons, img, confidence_threshold):
+    """Generate skeleton overlay on image (inplace).
+
+    Args:
+        skeletons (dict): skeleton joint co-ordinates
+        img (np.array): source image
+        confidence_threshold (float): Threshold
+    """
     skeleton_color = (100, 254, 213)
     for index, skeleton in enumerate(skeletons):
         limbs = get_valid_limbs(keypoint_ids, skeleton, confidence_threshold)
         for limb in limbs:
             cv2.line(
                 img, limb[0], limb[1], skeleton_color, thickness=2, lineType=cv2.LINE_AA
-            )
-
-# parser = argparse.ArgumentParser(description="Perform keypoing estimation on an image")
-# parser.add_argument(
-#     "-c",
-#     "--confidence_threshold",
-#     type=float,
-#     default=0.5,
-#     help="Minimum confidence (0-1) of displayed joints",
-# )
-# parser.add_argument(
-#     "-v",
-#     "--verbose",
-#     action="store_true",
-#     help="Increase output verbosity by enabling backend logging",
-# )
-
-# parser.add_argument(
-#     "-o",
-#     "--output_image",
-#     type=str,
-#     help="filename of the output image",
-# )
-
-# parser.add_argument("image", metavar="I", type=str, help="filename of the input image")
+            )

pose_tracking/cubemos_wrapper.py

@@ -1,6 +1,12 @@
-#!/usr/bin/env python3
-
-# cubemos init
+#!/usr/bin/env python
+# Title           :loader.py
+# Author          :Venkatraman Narayanan, Bala Murali Manoghar, Vishnu Shashank Dorbala, Aniket Bera, Dinesh Manocha
+# Copyright       :"Copyright 2020, Proxemo project"
+# Version         :1.0
+# License         :"MIT"
+# Maintainer      :Venkatraman Narayanan, Bala Murali Manoghar
+# Email           :[email protected], [email protected]
+# ==============================================================================
 import os
 import numpy as np
 import cv2
@@ -37,7 +43,15 @@
 
 
 class Cubemos_Tacker():
+    """Cubemos skeletal tracker class."""
+
     def __init__(self, intrinsics, verbose=False):
+        """Constructor.
+
+        Args:
+            intrinsics (np.array): Camera instrinsics matrix
+            verbose (bool, optional): Verbose log. Defaults to False.
+        """
         self.verbose = verbose
         self.intrinsics = intrinsics
         self.confidence_threshold = 0.3
@@ -54,6 +68,7 @@ def __init__(self, intrinsics, verbose=False):
         self.init_cubemos_api()
 
     def init_skel_track(self):
+        """Skeleton tracking initializer."""
         check_license_and_variables_exist()
         # Get the path of the native libraries and ressource files
         self.sdk_path = os.environ["CUBEMOS_SKEL_SDK"]
@@ -64,6 +79,7 @@ def init_skel_track(self):
                                default_log_dir())
 
     def init_cubemos_api(self):
+        """Setup cubemos API."""
         # initialize the api with a valid license key in default_license_dir()
         self.api = Api(default_license_dir())
         model_path = os.path.join(self.sdk_path,
@@ -74,6 +90,12 @@ def init_cubemos_api(self):
         self.api.load_model(CM_TargetComputeDevice.CM_CPU, model_path)
 
     def track_skeletons(self, image, depth_image):
+        """Track skeleton from image frame to deptg frame.
+
+        Args:
+            image (np.array): RGB image
+            depth_image (np.array): depth image
+        """
         # perform inference
         self.skeletons = self.api.estimate_keypoints(
             image, self.network_height)
@@ -88,6 +110,11 @@ def track_skeletons(self, image, depth_image):
         self.skel2D_to_skel3D(depth_image)
 
     def render_skeletons(self, image):
+        """Overlay detected skeleton on image.
+
+        Args:
+            image (np.array): RGB image
+        """
         render_result(self.skeletons, image, self.confidence_threshold)
         for skel_num, joints in enumerate(self.skel3d_np):
             for joint_ndx, pt_3D in enumerate(joints[[0, -1]]):
@@ -115,6 +142,15 @@ def render_skeletons(self, image):
                            2, (255, 0, 255), -1)
 
     def map_2D_3D(self, pixel, depth):
+        """Match camera frame to depth frame
+
+        Args:
+            pixel (np.array): pixel co-ordinates
+            depth (np.array): depth co-ordinates
+
+        Returns:
+            [np.array]: 3D point co-ordinates
+        """
         x = (pixel[0] - self.intrinsics.ppx) / self.intrinsics.fx
         y = (pixel[1] - self.intrinsics.ppy) / self.intrinsics.fy
 
@@ -139,6 +175,11 @@ def map_2D_3D(self, pixel, depth):
         return (X, Y, Z)
 
     def skel2D_to_skel3D(self, depth_image):
+        """Convert a 2D skeleton in Camera frame to 3D co-ordinates.
+
+        Args:
+            depth_image (np.array): depth image
+        """
         self.skel_ids = []
         self.skel3d_np = []
         self.roots_2d = []

pose_tracking/human_tracking_3D.py

@@ -1,5 +1,12 @@
-#!/usr/bin/env python3
-
+#!/usr/bin/env python
+# Title           :loader.py
+# Author          :Venkatraman Narayanan, Bala Murali Manoghar, Vishnu Shashank Dorbala, Aniket Bera, Dinesh Manocha
+# Copyright       :"Copyright 2020, Proxemo project"
+# Version         :1.0
+# License         :"MIT"
+# Maintainer      :Venkatraman Narayanan, Bala Murali Manoghar
+# Email           :[email protected], [email protected]
+# ==============================================================================
 import cv2
 import numpy as np
 
@@ -8,18 +15,39 @@
 
 
 class Skel_Temporal():
+    """Skeleton gait generator class."""
+
     def __init__(self, skel_id, do_not_ignore_false_limbs=True):
+        """Constructor
+
+        Args:
+            skel_id (int): Skeleton ID
+            do_not_ignore_false_limbs (bool, optional): Ignore false limbs?. Defaults to True.
+        """
         self.id = skel_id
         self.skel_temporal = []
         self.do_not_ignore_false_limbs = do_not_ignore_false_limbs
 
     def add(self, skel_ti):
+        """Add skeleton to tracking.
+
+        Args:
+            skel_ti (np.array): skeleton co-ordinates
+        """
         if not np.any(skel_ti == -1) or self.ignore_false_limbs:
             if len(self.skel_temporal) > 75:
                 self.skel_temporal.pop(0)
             self.skel_temporal.append(skel_ti)
 
     def __eq__(self, other):
+        """Skeleton compare
+
+        Args:
+            other (obj): Skeleton Object
+
+        Returns:
+            [bool]: Same/different skeleton
+        """
         try:
             if self.id == other.id:
                 return True
@@ -32,6 +60,11 @@ def __eq__(self, other):
                 return False
 
     def get_embedding(self):
+        """Convert Temporal gait cycle to Image sequence.
+
+        Returns:
+            [np.array]: Gait cycle embedded as image
+        """
         skel_temporal_np = np.array(self.skel_temporal)
         # make root as (0, 0, 0)
         # even if number of frames is less than 75 it will be resized to 244*244
@@ -42,12 +75,25 @@ def get_embedding(self):
 
 
 class Skel_Tracker():
+    """Skeleton Tracking Class."""
+
     def __init__(self, do_not_ignore_false_limbs=True):
+        """Constructor.
+
+        Args:
+            do_not_ignore_false_limbs (bool, optional): Ignore false limbs?. Defaults to True.
+        """
         self.skel_tracks = []
         self.img_embeddings = []
         self.do_not_ignore_false_limbs = do_not_ignore_false_limbs
 
     def update(self, skel_nps, skel_ids):
+        """Add skeleton pose to sequence.
+
+        Args:
+            skel_nps (np.array): Skeleton co-ordinates
+            skel_ids (list): Skeleton IDs
+        """
         # add skeleton corresponding to id
         for skel_np, skel_id in zip(skel_nps, skel_ids):
             try:
@@ -75,6 +121,11 @@ def update(self, skel_nps, skel_ids):
             self.skel_tracks.pop(value - ndx)
 
     def get_embedding(self):
+        """Generate image embedding for entire gait sequence.
+
+        Returns:
+            [list]: image embeddings, skeleton IDs
+        """
         self.img_embeddings = []
         ids = []
         for skel_track in self.skel_tracks:
@@ -84,6 +135,7 @@ def get_embedding(self):
         return self.img_embeddings, ids
 
     def display_embedding(self):
+        """View image embeddings."""
         imgs = self.img_embeddings[0]  # np.empty((244,244,3))
         print(self.img_embeddings.shape)
         for img in self.img_embeddings[1:]:
@@ -94,7 +146,15 @@ def display_embedding(self):
 
 
 class Track_Human_Pose():
+    """Main gait tracking loop."""
+
     def __init__(self, display=True, verbose=True):
+        """Constructor.
+
+        Args:
+            display (bool, optional): Show skeleton detections?. Defaults to True.
+            verbose (bool, optional): Generate verbose log?. Defaults to True.
+        """
         self.verbose = verbose
         self.display = display
 
@@ -104,6 +164,7 @@ def __init__(self, display=True, verbose=True):
         self.skel_tracker = Skel_Tracker()
 
     def get_pose(self):
+        """Get human skeletons."""
         # capture
         self.camera.capture()
         # get skeletons
@@ -121,10 +182,12 @@ def get_pose(self):
             cv2.imshow('RealSense', images)
 
     def track_pose(self):
+        """Track skeleton gaits."""
         self.skel_tracker.update(self.cubemos.skel3d_np,
                                  self.cubemos.skel_ids)
 
     def cleanup(self):
+        """Cleanup workspace setup."""
         self.camera.cleanup()