alobugdays/23-pad-accepts-multiple-argument

aloscene/tensors/augmented_tensor.py

@@ -962,7 +962,7 @@ def _pad_label(self, label, offset_y, offset_x, **kwargs):
         except AttributeError:
             return label
 
-    def pad(self, offset_y: tuple, offset_x: tuple, **kwargs):
+    def pad(self, offset_y: tuple = None, offset_x: tuple = None, multiple: int = None, **kwargs):
         """
         Pad AugmentedTensor, and its labels recursively
 
@@ -974,16 +974,34 @@ def pad(self, offset_y: tuple, offset_x: tuple, **kwargs):
         offset_x: tuple of float or tuple of int
             (percentage left_offset, percentage right_offset) Percentage based on the previous size. If tuple of int
             the absolute value will be converted to float (percentage) before to be applied.
+        multiple: int
+            pad the tensor to the next multiple of `multiple`
 
         Returns
         -------
-        croped : aloscene AugmentedTensor
-            croped tensor
+        cropped : aloscene AugmentedTensor
+            cropped tensor
         """
-        if isinstance(offset_y[0], int) and isinstance(offset_y[1], int):
-            offset_y = (offset_y[0] / self.H, offset_y[1] / self.H)
-        if isinstance(offset_x[0], int) and isinstance(offset_x[1], int):
-            offset_x = (offset_x[0] / self.W, offset_x[1] / self.W)
+        if multiple is not None:
+            assert offset_x is None and offset_y is None
+            if not self.H % multiple == 0:
+                offset_y0 = int(np.floor((multiple - self.H % multiple) / 2))
+                offset_y1 = int(np.ceil((multiple - self.H % multiple) / 2))
+                offset_y = (offset_y0 / self.H, offset_y1 / self.H)
+            else:  # already a multiple of H
+                offset_y = (0, 0)
+            if not self.W % multiple == 0:
+                offset_x0 = int(np.floor((multiple - self.W % multiple) / 2))
+                offset_x1 = int(np.ceil((multiple - self.W % multiple) / 2))
+                offset_x = (offset_x0 / self.W, offset_x1 / self.W)
+            else:  # already a multiple of W
+                offset_x = (0, 0)
+        else:
+            assert offset_x is not None and offset_y is not None
+            if isinstance(offset_y[0], int) and isinstance(offset_y[1], int):
+                offset_y = (offset_y[0] / self.H, offset_y[1] / self.H)
+            if isinstance(offset_x[0], int) and isinstance(offset_x[1], int):
+                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))

unittest/test_frame.py

@@ -451,7 +451,7 @@ def test_recursive_temporal_batch_bis():
 
 def test_pad():
     S = 600
-    frame = aloscene.Frame(np.random.uniform(0, 1, (3, 600, 600)), normalization="01", names=("C", "H", "W"))
+    frame = aloscene.Frame(np.random.uniform(0, 1, (3, S, S)), normalization="01", names=("C", "H", "W"))
     frame = frame.norm_resnet()
     boxes1 = BoundingBoxes2D(
         np.array([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.1, 0.1]]), boxes_format="xcyc", absolute=False, names=("N", None)
@@ -463,6 +463,11 @@ def test_pad():
     frame.append_boxes2d(boxes2, "boxes2")
     n_frame = frame.pad(offset_y=(0.0, 0.1), offset_x=(0.0, 0.1))
 
+    # test argument `multiple`
+    for m in range(1, 256):
+        new_h, new_w = frame.pad(multiple=m).shape[-2:]
+        assert (new_h % m == 0) and (new_w % m == 0)
+
 
 def test_device_propagation():