⚡️ Speed up method ResizeGenerator.forward by 99%

kornia/augmentation/random_generator/_2d/resize.py

@@ -77,11 +77,12 @@ def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) ->
         input_size = h, w = (batch_shape[-2], batch_shape[-1])
 
         src = bbox_generator(
-            torch.tensor(0, device=_device, dtype=_dtype),
-            torch.tensor(0, device=_device, dtype=_dtype),
-            torch.tensor(input_size[1], device=_device, dtype=_dtype),
-            torch.tensor(input_size[0], device=_device, dtype=_dtype),
-        ).repeat(batch_size, 1, 1)
+            torch.full((batch_size,), 0, device=_device, dtype=_dtype),
+            torch.full((batch_size,), 0, device=_device, dtype=_dtype),
+            torch.full((batch_size,), input_size[1], device=_device, dtype=_dtype),
+            torch.full((batch_size,), input_size[0], device=_device, dtype=_dtype),
+        )
+
 
         if isinstance(self.output_size, int):
             aspect_ratio = w / h
@@ -99,11 +100,12 @@ def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) ->
             raise AssertionError(f"`resize_to` must be a tuple of 2 positive integers. Got {output_size}.")
 
         dst = bbox_generator(
-            torch.tensor(0, device=_device, dtype=_dtype),
-            torch.tensor(0, device=_device, dtype=_dtype),
-            torch.tensor(output_size[1], device=_device, dtype=_dtype),
-            torch.tensor(output_size[0], device=_device, dtype=_dtype),
-        ).repeat(batch_size, 1, 1)
+            torch.full((batch_size,), 0, device=_device, dtype=_dtype),
+            torch.full((batch_size,), 0, device=_device, dtype=_dtype),
+            torch.full((batch_size,), output_size[1], device=_device, dtype=_dtype),
+            torch.full((batch_size,), output_size[0], device=_device, dtype=_dtype),
+        )
+
 
         _input_size = torch.tensor(input_size, device=_device, dtype=torch.long).expand(batch_size, -1)
         _output_size = torch.tensor(output_size, device=_device, dtype=torch.long).expand(batch_size, -1)

kornia/geometry/bbox.py

@@ -365,16 +365,21 @@ def bbox_generator(
             f"`width`({width.device}), `height`({height.device})."
         )
 
-    bbox = torch.tensor([[[0, 0], [0, 0], [0, 0], [0, 0]]], device=x_start.device, dtype=x_start.dtype).repeat(
-        1 if x_start.dim() == 0 else len(x_start), 1, 1
-    )
-
-    bbox[:, :, 0] += x_start.view(-1, 1)
-    bbox[:, :, 1] += y_start.view(-1, 1)
-    bbox[:, 1, 0] += width - 1
-    bbox[:, 2, 0] += width - 1
-    bbox[:, 2, 1] += height - 1
-    bbox[:, 3, 1] += height - 1
+    if x_start.dim() == 0:
+        x_start = x_start.unsqueeze(0)
+        y_start = y_start.unsqueeze(0)
+        width = width.unsqueeze(0)
+        height = height.unsqueeze(0)
+
+    x1 = x_start + width - 1
+    y1 = y_start + height - 1
+
+    bbox = torch.stack([
+        torch.stack([x_start, y_start], dim=1),
+        torch.stack([x1, y_start], dim=1),
+        torch.stack([x1, y1], dim=1),
+        torch.stack([x_start, y1], dim=1),
+    ], dim=1)
 
     return bbox
 

kornia/geometry/transform/affwarp.py

@@ -26,7 +26,8 @@
 from kornia.filters import gaussian_blur2d
 from kornia.image.utils import perform_keep_shape_image
 
-from .imgwarp import get_affine_matrix2d, get_projective_transform, get_rotation_matrix2d, warp_affine, warp_affine3d
+from .imgwarp import (get_affine_matrix2d, get_projective_transform,
+                      get_rotation_matrix2d, warp_affine, warp_affine3d)
 
 __all__ = [
     "Affine",