numpy and tensorflow docstrings

Author: rfahrn

src/python/pose_format/numpy/pose_body.py

@@ -12,9 +12,38 @@
 # np.seterr(all='raise')
 
 class NumPyPoseBody(PoseBody):
-    tensor_reader = 'unpack_numpy'
+    """Represents pose information leveraging NumPy operations and structures.
+
+     * Inherits from:  `PoseBody`
+     * Implements pose info using NumPy operations and structures.
+     * Provides method for operations:  matrix, multiplication, interpolation and data type conversions 
+
+    The `NumPyPoseBody` is an implementation of  `PoseBody` base class. 
+    This subclass uses NumPy masked arrays to handle pose data. 
+    Makes it suitable for applications where you need NumPy-based operations. 
+    The masked arrays allow for efficient handling of missing or invalid pose values
+
+    The class also comes with methods to transform, modify, and operate on pose data, including matrix multiplication, interpolation, 
+    and conversions to other data types like PyTorch tensors or TensorFlow tensors
+    
+    Parameters
+    ----------
+    fps : float
+        Frames per second, to represent the temporal aspect of pose data.
+    data : Union[ma.MaskedArray, np.ndarray]
+        Pose data either as a masked array or a regular numpy array.
+    confidence : np.ndarray
+        confidence array of the pose keypoints.
+    
+    """
+
+    tensor_reader = 'unpack_numpy' 
+    """Specifies the method name for unpacking a numpy array (Value: 'unpack_numpy')."""
 
     def __init__(self, fps: float, data: Union[ma.MaskedArray, np.ndarray], confidence: np.ndarray):
+        """
+        Initializes the NumPyPoseBody instance
+        """   
         if isinstance(data, np.ndarray):  # If array is not masked
             mask = confidence == 0  # 0 means no-mask, 1 means with-mask
             stacked_mask = np.stack([mask] * data.shape[-1], axis=3)
@@ -24,6 +53,21 @@ def __init__(self, fps: float, data: Union[ma.MaskedArray, np.ndarray], confiden
 
     @classmethod
     def read_v0_0(cls, header: PoseHeader, reader: BufferReader, **unused_kwargs):
+        """
+        Reads pose data from a given buffer reader using a specified data format version (see: ``docs/specs``).
+
+        Parameters
+        ----------
+        header : PoseHeader
+            Pose header information
+        reader : BufferReader
+            binary buffer reader
+
+        Returns
+        -------
+        NumPyPoseBody
+            Instance of NumPyPoseBody with read pose data.
+        """
         fps, _frames = reader.unpack(ConstStructs.double_ushort)
 
         _dims = max([len(c.format) for c in header.components]) - 1
@@ -64,6 +108,16 @@ def read_v0_0(cls, header: PoseHeader, reader: BufferReader, **unused_kwargs):
         return cls(fps, ma.stack(frames_d), ma.stack(frames_c))
 
     def write(self, version: float, buffer: BinaryIO):
+        """
+        Writes pose data to a binary buffer using specified data format version.
+
+        Parameters
+        ----------
+        version : float
+            Version of the data format.
+        buffer : BinaryIO
+            The binary buffer to write to.
+        """
         _frames, _people, _points, _dims = self.data.shape
         _frames = _frames if _frames < 65535 else 0  # TODO change from short to int
         buffer.write(ConstStructs.triple_ushort.pack(self.fps, _frames, _people))
@@ -73,9 +127,17 @@ def write(self, version: float, buffer: BinaryIO):
 
     @property
     def mask(self):
+        """ Returns  mask associated with data. """
         return self.data.mask
 
     def torch(self):
+        """converts current instance into a TorchPoseBody instance.
+
+        Returns
+        -------
+        TorchPoseBody
+            The pose body data represented in PyTorch tensors.
+        """
         try:
             import torch
         except ImportError:
@@ -91,6 +153,14 @@ def torch(self):
         return TorchPoseBody(self.fps, torch_data, torch_confidence)
 
     def tensorflow(self):
+        """
+        converts current instance into a TensorflowPoseBody instance
+
+        Returns
+        -------
+        TensorflowPoseBody
+            pose body data represented in TensorFlow tensors
+        """
         import tensorflow
         from ..tensorflow.pose_body import TensorflowPoseBody
 
@@ -99,24 +169,79 @@ def tensorflow(self):
         return TensorflowPoseBody(self.fps, tf_data, tf_confidence)
 
     def zero_filled(self):
+        """
+        fills missing values with zeros.
+
+        Returns
+        -------
+        NumPyPoseBody
+            changed pose body data.
+        """
         self.data = ma.array(self.data.filled(0), mask=self.data.mask)
         return self
 
     def matmul(self, matrix: np.ndarray):
+        """
+        Performs matrix multiplication on pose data.
+
+        Parameters
+        ----------
+        matrix : np.ndarray
+            matrix to multiply the pose data with
+
+        Returns
+        -------
+        NumPyPoseBody
+            transformed pose body data
+        """
         data = ma.dot(self.data, matrix)
         return NumPyPoseBody(self.fps, data, self.confidence)
 
     def flip(self, axis=0):
+        """
+        flips pose data across a specified axis
+
+        Parameters
+        ----------
+        axis : int, optional
+            axis along which the pose data should be flipped.
+
+        Returns
+        -------
+        NumPyPoseBody
+            flipped pose body data
+        """
         vec = np.ones(self.data.shape[-1])
         vec[axis] = -1
 
         data = self.data * vec
         return NumPyPoseBody(self.fps, data, self.confidence)
 
     def points_perspective(self):
+        """
+        Transforms pose data to get a perspective based on points.
+
+        Returns
+        -------
+        ma.MaskedArray
+            Transformed pose data
+        """
         return ma.transpose(self.data, axes=POINTS_DIMS)
 
     def get_points(self, indexes: List[int]):
+        """
+        Get points (keypoints) based on given indexes.
+
+        Parameters
+        ----------
+        indexes : List[int]
+             List of indices representing the keypoints to get.
+
+        Returns
+        -------
+        NumPyPoseBody
+            Pose body data containing only a specified points.
+        """
         data = ma.transpose(self.data, axes=POINTS_DIMS)
         new_data = ma.transpose(data[indexes], axes=POINTS_DIMS)
 
@@ -127,6 +252,19 @@ def get_points(self, indexes: List[int]):
         return NumPyPoseBody(self.fps, new_data, new_confidence)
 
     def bbox(self, header: PoseHeader):
+        """
+        Computes the bounding boxes for each component based on the pose data.
+
+        Parameters
+        ----------
+        header : PoseHeader
+            Pose header information.
+
+        Returns
+        -------
+        NumPyPoseBody
+            Pose body data representing bounding boxes.
+        """
         data = ma.transpose(self.data, axes=POINTS_DIMS)
 
         # Split data by components, `ma` doesn't support ".split"
@@ -151,6 +289,21 @@ def bbox(self, header: PoseHeader):
         return NumPyPoseBody(self.fps, new_data, confidence)
 
     def interpolate(self, new_fps: int = None, kind='cubic'):
+        """
+        Interpolates the pose data to match a new frame rate.
+
+        Parameters
+        ----------
+        new_fps : int, optional
+            The desired frame rate for interpolation.
+        kind : str, optional
+            The type of interpolation. Options include: "linear", "quadratic", and "cubic".
+
+        Returns
+        -------
+        NumPyPoseBody
+            Interpolated pose body data.
+        """
         try:
             from scipy.interpolate import interp1d
         except ImportError:
@@ -222,6 +375,17 @@ def interpolate(self, new_fps: int = None, kind='cubic'):
         return NumPyPoseBody(fps=new_fps, data=dimensions, confidence=confidence)
 
     def flatten(self):
+        """Flattens data and confidence arrays.
+
+        method reshapes data and confidence arrays to a two-dimensional array.
+        The flattened array is filtered to remove rows where confidence is zero.
+
+        Returns
+        --------
+        numpy.ndarray
+            flattened and filtered version of the data array.
+
+        """
         shape = self.data.shape
         data = self.data.data.reshape(-1, shape[-1])  # Not masked data
         confidence = self.confidence.flatten()

src/python/pose_format/numpy/representation/distance.py

@@ -2,7 +2,28 @@
 
 
 class DistanceRepresentation:
+    """A class to compute the Euclidean distance between two sets of points.
+    """
     def distance(self, p1s: ma.MaskedArray, p2s: ma.MaskedArray) -> ma.MaskedArray:
+        """
+        Compute the Euclidean distance between two sets of points.
+
+        Parameters
+        ----------
+        p1s : ma.MaskedArray
+            First set of points.
+        p2s : ma.MaskedArray
+            Second set of points.
+
+        Returns
+        -------
+        ma.MaskedArray
+            Euclidean distances between the two sets of points. The returned array has one fewer dimension than the input arrays, as the distance calculation collapses the last dimension.
+
+        Note
+        -----
+        this method assumes that input arrays `p1s` and `p2s` have same shape.
+        """
         diff = p1s - p2s
         square = ma.power(diff, 2)
         sum_squares = square.sum(axis=-1)
@@ -11,9 +32,18 @@ def distance(self, p1s: ma.MaskedArray, p2s: ma.MaskedArray) -> ma.MaskedArray:
 
     def __call__(self, p1s: ma.MaskedArray, p2s: ma.MaskedArray) -> ma.MaskedArray:
         """
-        Euclidean distance between two points
-        :param p1s: ma.MaskedArray (Points, Batch, Len, Dims)
-        :param p2s: ma.MaskedArray (Points, Batch, Len, Dims)
-        :return: ma.MaskedArray (Points, Batch, Len)
+        For `distance` method to compute Euclidean distance between two points.
+
+        Parameters
+        ----------
+        p1s : ma.MaskedArray, shape (Points, Batch, Len, Dims)
+            First set of points.
+        p2s : ma.MaskedArray, shape (Points, Batch, Len, Dims)
+            Second set of points.
+
+        Returns
+        -------
+        ma.MaskedArray, shape (Points, Batch, Len)
+            Euclidean distances between the two sets of points.
         """
         return self.distance(p1s, p2s)

src/python/pose_format/numpy/representation/distance_test.py

@@ -10,13 +10,35 @@
 
 
 class TestDistanceRepresentation(TestCase):
+    """
+    Unit tests for DistanceRepresentation class to computes Euclidean distance between sets of 3D points.
+    """
     def test_call_value_should_be_distance(self):
+        """
+        Test if computed distance between two points is correct.
+
+        Note
+        -----
+        This test initializes two sets of 3D points:
+        p1 = (1, 2, 3) and p2 = (4, 5, 6). It then checks if the 
+        computed Euclidean distance between p1 and p2 is sqrt(27).
+        """
+        
         p1s = ma.array([[[[1, 2, 3]]]], dtype=np.float32)
         p2s = ma.array([[[[4, 5, 6]]]], dtype=np.float32)
         distances = representation(p1s, p2s)
         self.assertAlmostEqual(float(distances[0][0][0]), math.sqrt(27), places=6)
 
     def test_call_masked_value_should_be_zero(self):
+        """
+        Test if the distance for masked values is zero.
+
+        Note
+        -----
+        This test checks scenario where one set of points 
+        is masked.The computed distance in such in such a case be 0.
+        """
+        
         mask = [[[[True, True, True]]]]
         p1s = ma.array([[[[1, 2, 3]]]], mask=mask, dtype=np.float32)
         p2s = ma.array([[[[4, 5, 6]]]], dtype=np.float32)