Feature/pose copy and remove components (#148)

* CDL: minor doc typo fix

* Undoing some changes that got mixed in

* Add Pose .copy() and .remove_components()

* fix type annotation in pose remove_components

* Adding copy functions to posebodies, and tests for this

* Some pylint changes

* Fix return type annotations, use copy() in zero_filled

* uncomment pytests for torchposebody, rename MaskedTensor imports

* import numpy.ma instead of 'from numpy import ma'

Author: cleong110

src/python/pose_format/numpy/pose_body.py

@@ -128,6 +128,11 @@ def write(self, version: float, buffer: BinaryIO):
         buffer.write(np.array(self.data.data, dtype=np.float32).tobytes())
         buffer.write(np.array(self.confidence, dtype=np.float32).tobytes())
 
+    def copy(self) -> 'NumPyPoseBody':
+        return type(self)(fps=self.fps,
+                          data=self.data.copy(),
+                          confidence=self.confidence.copy())
+
     @property
     def mask(self):
         """ Returns  mask associated with data. """
@@ -181,8 +186,9 @@ def zero_filled(self):
         NumPyPoseBody
             changed pose body data.
         """
-        self.data = ma.array(self.data.filled(0), mask=self.data.mask)
-        return self
+        copy = self.copy()
+        copy.data = ma.array(copy.data.filled(0), mask=copy.data.mask)
+        return copy
 
     def matmul(self, matrix: np.ndarray):
         """

src/python/pose_format/pose.py

@@ -202,6 +202,28 @@ def frame_dropout_normal(self, dropout_mean: float = 0.5, dropout_std: float = 0
         """
         body, selected_indexes = self.body.frame_dropout_normal(dropout_mean=dropout_mean, dropout_std=dropout_std)
         return Pose(header=self.header, body=body), selected_indexes
+    
+    
+    def remove_components(self, components_to_remove: Union[str, List[str]], points_to_remove: Union[Dict[str, List[str]],None] = None):
+        
+        if isinstance(components_to_remove, str):
+            components_to_remove = [components_to_remove]
+
+        components_to_keep = []
+        points_dict = {}
+
+        for component in self.header.components:
+            if component.name not in components_to_remove:
+                components_to_keep.append(component.name)
+                points_dict[component.name] = []
+                if points_to_remove is not None:
+                    for point in component.points:                    
+                        if point not in points_to_remove[component.name]:
+                            points_dict[component.name].append(point)
+
+        return self.get_components(components_to_keep, points_dict)
+        
+    
 
     def get_components(self, components: List[str], points: Union[Dict[str, List[str]],None] = None):
         """
@@ -253,6 +275,10 @@ def get_components(self, components: List[str], points: Union[Dict[str, List[str
         new_body = self.body.get_points(flat_indexes)
 
         return Pose(header=new_header, body=new_body)
+    
+
+    def copy(self):
+        return self.__class__(self.header, self.body.copy())
 
     def bbox(self):
         """

src/python/pose_format/pose_body.py

@@ -1,8 +1,9 @@
+import math
 from random import sample
-from typing import BinaryIO, List, Tuple
+from typing import BinaryIO, List, Tuple, Optional
 
 import numpy as np
-import math
+
 
 from pose_format.pose_header import PoseHeader
 from pose_format.utils.reader import BufferReader, ConstStructs
@@ -60,9 +61,9 @@ def read(cls, header: PoseHeader, reader: BufferReader, **kwargs) -> "PoseBody":
 
         if header.version == 0:
             return cls.read_v0_0(header, reader, **kwargs)
-        elif round(header.version, 3) == 0.1:
+        if round(header.version, 3) == 0.1:
             return cls.read_v0_1(header, reader, **kwargs)
-        elif round(header.version, 3) == 0.2:
+        if round(header.version, 3) == 0.2:
             return cls.read_v0_2(header, reader, **kwargs)
 
         raise NotImplementedError("Unknown version - %f" % header.version)
@@ -93,8 +94,8 @@ def read_v0_1_frames(cls,
                          frames: int,
                          shape: List[int],
                          reader: BufferReader,
-                         start_frame: int = None,
-                         end_frame: int = None):
+                         start_frame: Optional[int] = None,
+                         end_frame: Optional[int] = None):
         """
         Reads frame data for version 0.1 from a buffer.
 
@@ -149,8 +150,8 @@ def read_v0_1_frames(cls,
     def read_v0_1(cls,
                   header: PoseHeader,
                   reader: BufferReader,
-                  start_frame: int = None,
-                  end_frame: int = None,
+                  start_frame: Optional[int] = None,
+                  end_frame: Optional[int] = None,
                   **unused_kwargs) -> "PoseBody":
         """
         Reads pose data for version 0.1 from a buffer.
@@ -176,7 +177,7 @@ def read_v0_1(cls,
         fps, _frames = reader.unpack(ConstStructs.double_ushort)
 
         _people = reader.unpack(ConstStructs.ushort)
-        _points = sum([len(c.points) for c in header.components])
+        _points = sum(len(c.points) for c in header.components)
         _dims = header.num_dims()
 
         # _frames is defined as short, which sometimes is not enough! TODO change to int
@@ -191,10 +192,10 @@ def read_v0_1(cls,
     def read_v0_2(cls,
                   header: PoseHeader,
                   reader: BufferReader,
-                  start_frame: int = None,
-                  end_frame: int = None,
-                  start_time: int = None,
-                  end_time: int = None,
+                  start_frame: Optional[int] = None,
+                  end_frame: Optional[int] = None,
+                  start_time: Optional[int] = None,
+                  end_time: Optional[int] = None,
                   **unused_kwargs) -> "PoseBody":
         """
         Reads pose data for version 0.2 from a buffer.
@@ -256,6 +257,11 @@ def write(self, version: float, buffer: BinaryIO):
             Buffer to write the pose data to.
         """
         raise NotImplementedError("'write' not implemented on '%s'" % self.__class__)
+    
+    def copy(self)->"PoseBody":
+        return self.__class__(fps=self.fps,
+                          data=self.data,
+                          confidence=self.confidence)
 
     def __getitem__(self, index):
         """
@@ -306,7 +312,7 @@ def torch(self):
         Raises
         ------
         NotImplementedError
-            If toch is not implemented.
+            If torch is not implemented.
         """
         raise NotImplementedError("'torch' not implemented on '%s'" % self.__class__)
 
@@ -474,7 +480,7 @@ def get_points(self, indexes: List[int]) -> __qualname__:
         Returns
         -------
         PoseBody
-            PoseBody instance containing only choosen points.
+            PoseBody instance containing only chosen points.
              
         Raises
         ------

src/python/pose_format/tensorflow/pose_body.py

@@ -17,7 +17,7 @@ class TensorflowPoseBody(PoseBody):
     """
     Representation of pose body data, optimized for TensorFlow operations.
 
-    * Inherites from PoseBody 
+    * Inherits from PoseBody 
 
     Parameters
     ----------
@@ -43,10 +43,11 @@ def __init__(self, fps: float, data: Union[MaskedTensor, tf.Tensor], confidence:
 
         super().__init__(fps, data, confidence)
 
-    def zero_filled(self):
+    def zero_filled(self) -> 'TensorflowPoseBody':
         """Return an instance with zero-filled data."""
-        self.data = self.data.zero_filled()
-        return self
+        copy = self.copy()
+        copy.data = self.data.zero_filled()
+        return copy
 
     def select_frames(self, frame_indexes: List[int]):
         """
@@ -152,6 +153,17 @@ def points_perspective(self) -> MaskedTensor:
         """
         return self.data.transpose(perm=POINTS_DIMS)
 
+    def copy(self) -> 'TensorflowPoseBody':
+        # Ensure copies are fully detached from the TF computation graph by round-trip through numpy
+        detached_data = tf.convert_to_tensor(self.data.tensor.numpy())
+        detached_mask = tf.convert_to_tensor(self.data.mask.numpy())
+        data_copy = MaskedTensor(detached_data, detached_mask)
+        confidence_copy = tf.convert_to_tensor(self.confidence.numpy())
+        return self.__class__(
+            fps=self.fps,
+            data=data_copy,
+            confidence=confidence_copy)
+
     def get_points(self, indexes: List[int]):
         """
         Gets and returns points from pose data based on indexes 

src/python/pose_format/torch/pose_body.py

@@ -4,8 +4,6 @@
 import torch
 
 from ..pose_body import POINTS_DIMS, PoseBody
-from ..pose_header import PoseHeader
-from ..utils.reader import BufferReader
 from .masked.tensor import MaskedTensor
 
 
@@ -28,10 +26,21 @@ def __init__(self, fps: float, data: Union[MaskedTensor, torch.Tensor], confiden
         super().__init__(fps, data, confidence)
 
     def cuda(self):
-        """Move data and cofidence of tensors to GPU"""
+        """Move data and confidence of tensors to GPU"""
         self.data = self.data.cuda()
         self.confidence = self.confidence.cuda()
 
+    def copy(self) -> 'TorchPoseBody':
+        data_copy = MaskedTensor(tensor=self.data.tensor.detach().clone().to(self.data.tensor.device),
+                                 mask=self.data.mask.detach().clone().to(self.data.mask.device),
+                                 )
+        confidence_copy = self.confidence.detach().clone().to(self.confidence.device)
+
+        return self.__class__(fps=self.fps,
+                             data=data_copy,
+                             confidence=confidence_copy)
+
+
     def zero_filled(self) -> 'TorchPoseBody':
         """
         Fill invalid values with zeros.
@@ -42,8 +51,9 @@ def zero_filled(self) -> 'TorchPoseBody':
             TorchPoseBody instance with masked data filled with zeros.
 
         """
-        self.data.zero_filled()
-        return self
+        copy = self.copy()
+        copy.data = copy.data.zero_filled()
+        return copy
 
     def matmul(self, matrix: np.ndarray) -> 'TorchPoseBody':
         """
@@ -120,3 +130,6 @@ def flatten(self):
         scalar = torch.ones(len(shape) + shape[-1], device=data.device)
         scalar[0] = 1 / self.fps
         return flat * scalar
+
+
+