add_rotation

alodataset/transforms.py

@@ -554,6 +554,38 @@ def apply(self, frame: Frame):
         return frame
 
 
+class Rotate(AloTransform):
+    def __init__(self, angle: float, *args, **kwargs):
+        """Rotate the given frame using the given rotation angle.
+
+        Parameters
+        ----------
+        angle: float, between 0 and 360
+        """
+        assert isinstance(angle, float)
+        self.angle = angle
+        super().__init__(*args, **kwargs)
+
+    def sample_params(self):
+        """Sample an `angle` from the list of possible `angles`"""
+        return (self.angle,)
+
+    def set_params(self, angle):
+        """Given predefined params, set the params on the class"""
+        self.angle = angle
+
+    def apply(self, frame: Frame):
+        """Apply the transformation
+
+        Parameters
+        ----------
+        frame: Frame
+            Frame to apply the transformation on
+        """
+        frame = frame.rotate(self.angle)
+        return frame
+
+
 class RealisticNoise(AloTransform):
     def __init__(self, gaussian_std: float = 0.02, shot_std: float = 0.2, *args, **kwargs):
         """Add an approximation of a realistic noise to the image.

aloscene/points_2d.py

@@ -429,6 +429,42 @@ def _resize(self, size, **kwargs):
             abs_size = tuple(s * fs for s, fs in zip(size, points.frame_size))
             return points.abs_pos(abs_size)
 
+    def _rotate(self, angle, **kwargs):
+        """Rotate Point2d, but not their labels
+
+        Parameters
+        ----------
+        size : float
+
+        Returns
+        -------
+        points : aloscene.Point2d
+            rotated points
+        """
+        points = self.xy().clone()
+        H, W = self.frame_size
+        angle_rad = angle * np.pi / 180
+        rot_mat = torch.tensor([[np.cos(angle_rad), np.sin(angle_rad)], [-np.sin(angle_rad), np.cos(angle_rad)]]).to(
+            torch.float32
+        )
+        tr_mat = torch.tensor([W / 2, H / 2])
+        for i in range(points.shape[0]):
+            points[i] = torch.matmul(rot_mat, points[i] - tr_mat) + tr_mat
+
+        max_size = torch.as_tensor([W, H], dtype=torch.float32)
+        points_filter = (points >= 0).as_tensor() & (points <= max_size).as_tensor()
+        points_filter = points_filter[:, 0] & points_filter[:, 1]
+        points = points[points_filter]
+
+        points = points.rel_pos()
+        # no modification needed for relative coordinates
+        if self.absolute:
+            points = points.abs_pos(self.frame_size)
+
+        points = points.get_with_format(self.points_format)
+
+        return points
+
     def _crop(self, H_crop: tuple, W_crop: tuple, **kwargs):
         """Crop Points with the given relative crop
 
@@ -580,7 +616,44 @@ def _pad(self, offset_y: tuple, offset_x: tuple, pad_points2d: bool = True, **kw
         return points
 
     def _spatial_shift(self, shift_y: float, shift_x: float, **kwargs):
-        raise Exception("Not handle by points 2D")
+        """
+        Spatially shift the Points
+        Parameters
+        ----------
+        shift_y: float
+            Shift percentage on the y axis. Could be negative or positive
+        shift_x: float
+            Shift percentage on the x axis. Could ne negative or positive.
+        Returns
+        -------
+        shifted_tensor: aloscene.AugmentedTensor
+            shifted tensor
+        """
+        if self.padded_size is not None:
+            raise Exception(
+                "Can't process spatial shift when padded size is not Note. Call fit_to_padded_size() first"
+            )
+        # get needed information
+        original_format = self.points_format
+        original_absolute = self.absolute
+        frame_size = self.frame_size
+
+        # shift the points
+        n_points = self.clone().rel_pos().xy()
+        n_points += torch.as_tensor([[shift_x, shift_y]])  # , device=self.device)
+
+        # filter points outside of the image
+        points_filter = (n_points >= 0).as_tensor() & (n_points <= 1).as_tensor()
+        points_filter = points_filter[:, 0] & points_filter[:, 1]
+        n_points = n_points[points_filter]
+        n_points = n_points.reset_names()
+
+        # Put back the instance into the same state as before
+        if original_absolute:
+            n_points = n_points.abs_pos(frame_size)
+        n_points = n_points.get_with_format(original_format)
+
+        return n_points
 
     def as_points(self, points):
         n_points = self.clone()

aloscene/tensors/augmented_tensor.py

@@ -48,34 +48,29 @@ def __init__(self, x, **kwargs):
         super().__init__()
 
     def drop_children(self):
-        """ Remove all children from this augmented tensor and return
+        """Remove all children from this augmented tensor and return
         the removed children.
         """
         labels = {}
         for name in self._children_list:
-            labels[name] = {
-                "value": getattr(self, name),
-                "property": self._child_property[name]
-            }
+            labels[name] = {"value": getattr(self, name), "property": self._child_property[name]}
             setattr(self, name, None)
         return labels
 
     def get_children(self):
-        """ Get all children attached from this augmented tensor
-        """
+        """Get all children attached from this augmented tensor"""
         labels = {}
         for name in self._children_list:
             # Use apply to return the same label but with a new structure
             # so that if the returned structure is changed, this will not impact the current one
             labels[name] = {
                 "value": self.apply_on_child(getattr(self, name), lambda l: l),
-                "property": self._child_property[name]
+                "property": self._child_property[name],
             }
         return labels
 
     def set_children(self, labels):
-        """ Set children on this augmented tensor
-        """
+        """Set children on this augmented tensor"""
         for name in labels:
             if name not in self._children_list:
                 self.add_child(name, labels[name]["value"], **labels[name]["property"])
@@ -153,7 +148,11 @@ def __check_child_name_alignment(var):
             elif not any(v in self.COMMON_DIM_NAMES for v in var.names):
                 return True
             for dim_id, dim_name in enumerate(var.names):
-                if dim_id < len(self.names) and dim_name in self.COMMON_DIM_NAMES and self.names[dim_id] == dim_name:
+                if (
+                    dim_id < len(self.names)
+                    and dim_name in self.COMMON_DIM_NAMES
+                    and self.names[dim_id] == dim_name
+                ):
                     return True
             raise Exception(
                 f"Impossible to align label name dim ({var.names}) with the tensor name dim ({self.names})."
@@ -166,7 +165,7 @@ def __check_child_name_alignment(var):
         return True
 
     def add_child(self, child_name, child, align_dim=["B", "T"], mergeable=True, **kwargs):
-        """ Add/attached an augmented tensor on this augmented tensor.
+        """Add/attached an augmented tensor on this augmented tensor.
 
         Parameters
         ----------
@@ -198,8 +197,6 @@ def add_child(self, child_name, child, align_dim=["B", "T"], mergeable=True, **k
 
         setattr(self, child_name, child)
 
-
-
     def _append_child(self, child_name: str, child, set_name: str = None):
         """
         Attach a new value for a given child name.
@@ -238,8 +235,6 @@ def _append_child(self, child_name: str, child, set_name: str = None):
         else:
             label[set_name] = child
 
-
-
     def _getitem_child(self, label, label_name, idx):
         """
         This method is used in AugmentedTensor.__getitem__
@@ -398,7 +393,9 @@ def cuda(self, *args, **kwargs):
         for name in self._children_list:
             label = getattr(self, name)
             if label is not None:
-                setattr(n_frame, name, self.apply_on_child(label, lambda l: l.cuda(*args, **kwargs)))
+                setattr(
+                    n_frame, name, self.apply_on_child(label, lambda l: l.cuda(*args, **kwargs))
+                )
         return n_frame
 
     def _merge_child(self, label, label_name, key, dict_merge, kwargs, check_dim=True):
@@ -410,7 +407,10 @@ def _merge_child(self, label, label_name, key, dict_merge, kwargs, check_dim=Tru
             # merge everything on the real target dimension.
             target_dim = 0
 
-        if check_dim and self.names[target_dim] not in self._child_property[label_name]["align_dim"]:
+        if (
+            check_dim
+            and self.names[target_dim] not in self._child_property[label_name]["align_dim"]
+        ):
             raise Exception(
                 "Can only merge labeled tensor on the following dimension '{}'. \
                 \nDrop the labels before to apply such operations or convert your labeled tensor to tensor first.".format(
@@ -479,15 +479,21 @@ def _merge_tensor(self, n_tensor, tensor_list, func, types, args=(), kwargs=None
                     if isinstance(label_value, dict):
                         for key in label_value:
                             labels_dict2list[label_name] = self._merge_child(
-                                label_value[key], label_name, key, labels_dict2list[label_name], kwargs
+                                label_value[key],
+                                label_name,
+                                key,
+                                labels_dict2list[label_name],
+                                kwargs,
                             )
                     elif label_value is None and isinstance(labels_dict2list[label_name], dict):
                         for key in labels_dict2list[label_name]:
                             labels_dict2list[label_name] = self._merge_child(
                                 label_value, label_name, key, labels_dict2list[label_name], kwargs
                             )
                     else:
-                        self._merge_child(label_value, label_name, label_name, labels_dict2list, kwargs)
+                        self._merge_child(
+                            label_value, label_name, label_name, labels_dict2list, kwargs
+                        )
             else:
                 raise Exception("Can't merge none AugmentedTensor with AugmentedTensor")
 
@@ -532,7 +538,9 @@ def _handle_expand_on_label(label, name):
         for name in self._children_list:
             label = getattr(tensor, name)
             if label is not None:
-                results = self.apply_on_child(label, lambda l: _handle_expand_on_label(l, name), on_list=False)
+                results = self.apply_on_child(
+                    label, lambda l: _handle_expand_on_label(l, name), on_list=False
+                )
                 setattr(tensor, name, results)
 
     def __iter__(self):
@@ -713,7 +721,12 @@ def __repr__(self):
                     if isinstance(values[key], list):
                         cvalue = (
                             f"{key}:["
-                            + ", ".join(["{}".format(len(k) if k is not None else None) for k in values[key]])
+                            + ", ".join(
+                                [
+                                    "{}".format(len(k) if k is not None else None)
+                                    for k in values[key]
+                                ]
+                            )
                             + "]"
                         )
                         content_value += f"{cvalue}, "
@@ -766,6 +779,7 @@ def recursive_apply_on_children_(self, func):
         """
         Recursively apply function on labels to modify tensor inplace
         """
+
         def __apply(l):
             return func(l).recursive_apply_on_children_(func)
 
@@ -853,6 +867,32 @@ def resize_func(label):
     def _resize(self, *args, **kwargs):
         raise Exception("This Augmented tensor should implement this method")
 
+    def rotate(self, angle, **kwargs):
+        """
+        Rotate AugmentedTensor, and its labels recursively
+
+        Parameters
+        ----------
+        angle : float
+
+        Returns
+        -------
+        rotated : aloscene AugmentedTensor
+            rotated tensor
+        """
+
+        def rotate_func(label):
+            try:
+                label_rotated = label._rotate(angle, **kwargs)
+                return label_rotated
+            except AttributeError:
+                return label
+
+        rotated = self._rotate(angle, **kwargs)
+        rotated.recursive_apply_on_children_(rotate_func)
+
+        return rotated
+
     def _crop_label(self, label, H_crop, W_crop, **kwargs):
         try:
             label_resized = label._crop(H_crop, W_crop, **kwargs)
@@ -919,10 +959,11 @@ def pad(self, offset_y: tuple, offset_x: tuple, **kwargs):
             offset_x = (offset_x[0] / self.W, offset_x[1] / self.W)
 
         padded = self._pad(offset_y, offset_x, **kwargs)
-        padded.recursive_apply_on_children_(lambda label: self._pad_label(label, offset_y, offset_x, **kwargs))
+        padded.recursive_apply_on_children_(
+            lambda label: self._pad_label(label, offset_y, offset_x, **kwargs)
+        )
         return padded
 
-
     def _spatial_shift_label(self, label, shift_y, shift_x, **kwargs):
         try:
             label_shift = label._spatial_shift(shift_y, shift_x, **kwargs)
@@ -947,7 +988,9 @@ def spatial_shift(self, shift_y: float, shift_x: float, **kwargs):
             shifted tensor
         """
         shifted = self._spatial_shift(shift_y, shift_x, **kwargs)
-        shifted.recursive_apply_on_children_(lambda label: self._spatial_shift_label(label, shift_y, shift_x, **kwargs))
+        shifted.recursive_apply_on_children_(
+            lambda label: self._spatial_shift_label(label, shift_y, shift_x, **kwargs)
+        )
         return shifted
 
     def _spatial_shift(self, shift_y, shift_x, **kwargs):
@@ -980,6 +1023,10 @@ def _resize(self, *args, **kwargs):
         # Must be implement by child class to handle resize
         return self.clone()
 
+    def _rotate(self, *args, **kwargs):
+        # Must be implement by child class to handle rotate
+        return self.clone()
+
     def _crop(self, *args, **kwargs):
         # Must be implement by child class to handle crop
         return self.clone()

aloscene/tensors/spatial_augmented_tensor.py

@@ -457,6 +457,24 @@ def _resize(self, size, **kwargs):
             return self.rename(None).view(shapes).reset_names()
         return F.resize(self.rename(None), (h, w)).reset_names()
 
+    def _rotate(self, angle, **kwargs):
+        """Rotate SpatialAugmentedTensor, but not its labels
+
+        Parameters
+        ----------
+        angle : float
+
+        Returns
+        -------
+        rotated : aloscene.SpatialAugmentedTensor
+            rotated tensor
+        """
+        # If a SpatialAgumentedTensor is empty, rotate operation does not work. Use view instead.
+        assert not (
+            ("N" in self.names and self.size("N") == 0) or ("C" in self.names and self.size("C") == 0)
+        ), "rotation is not possible on an empty tensor"
+        return F.rotate(self.rename(None), angle).reset_names()
+
     def _crop(self, H_crop: tuple, W_crop: tuple, **kwargs):
         """Crop the SpatialAugmentedTensor