Repository: kornia/kornia
Pull request: [fix] RandomCrop fix and improvement
Fixes #1569 

1. fix `keepdim` bug
2. improve `same_on_batch` for slicing in a loop

Problem statement:
keepdim=True in RandomCrop not working
### Describe the bug

Setting `keepdim=True` in `RandomCrop()` has no effect. The batch dim is added either way.

### Reproduction steps

```bash
1. Initialize RandomCrop with keepdim=True
2. Transform a C,H,W tensor
```


### Expected behavior

Expected: Receive 3D tensor
Actual: Receive 4D tensor

### Environment

```shell
- PyTorch Version (e.g., 1.0): 1.10.0
- OS (e.g., Linux): Google Colab
- How you installed PyTorch (`conda`, `pip`, source): was installed
- Build command you used (if compiling from source): Nope
- Python version: 3.7.12
- CUDA/cuDNN version: 11.1
- GPU models and configuration: multiple 
- Any other relevant information: kornia 0.6.3
```


### Additional context

_No response_

Hints from issue discussion:
Thanks for reporting. Fixed in #1571

Commit message: Update kornia/augmentation/_2d/geometric/crop.py

Patch to review:
diff --git a/kornia/augmentation/_2d/geometric/center_crop.py b/kornia/augmentation/_2d/geometric/center_crop.py
index cf524b33d6..1693708575 100644
--- a/kornia/augmentation/_2d/geometric/center_crop.py
+++ b/kornia/augmentation/_2d/geometric/center_crop.py
@@ -110,14 +110,17 @@ def apply_transform(
         if self.flags["cropping_mode"] == "slice":  # uses advanced slicing to crop
             # TODO: implement as separated function `crop_and_resize_iterative`
             B, C, _, _ = input.shape
-            H, W = self.size
-            out = torch.empty(B, C, H, W, device=input.device, dtype=input.dtype)
+            x1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 0].numpy()
+            x2 = params["src"].to(device="cpu", dtype=torch.long)[:, 1, 0].numpy() + 1
+            y1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 1].numpy()
+            y2 = params["src"].to(device="cpu", dtype=torch.long)[:, 3, 1].numpy() + 1
+
+            if self.same_on_batch:
+                return input[..., y1[0]:y2[0], x1[0]:x2[0]]
+
+            out = torch.empty(B, C, *self.flags["size"], device=input.device, dtype=input.dtype)
             for i in range(B):
-                x1 = int(params["src"][i, 0, 0])
-                x2 = int(params["src"][i, 1, 0]) + 1
-                y1 = int(params["src"][i, 0, 1])
-                y2 = int(params["src"][i, 3, 1]) + 1
-                out[i] = input[i : i + 1, :, y1:y2, x1:x2]
+                out[i] = input[i : i + 1, :, y1[i]:y2[i], x1[i]:x2[i]]
             return out
         raise NotImplementedError(f"Not supported type: {self.flags['cropping_mode']}.")
 
diff --git a/kornia/augmentation/_2d/geometric/crop.py b/kornia/augmentation/_2d/geometric/crop.py
index 4e94bf419c..2ece32d3f6 100644
--- a/kornia/augmentation/_2d/geometric/crop.py
+++ b/kornia/augmentation/_2d/geometric/crop.py
@@ -6,7 +6,7 @@
 from kornia.augmentation import random_generator as rg
 from kornia.augmentation._2d.geometric.base import GeometricAugmentationBase2D
 from kornia.augmentation.base import TensorWithTransformMat
-from kornia.augmentation.utils import _transform_input
+from kornia.augmentation.utils import _transform_input, _transform_output_shape
 from kornia.constants import Resample
 from kornia.geometry.transform import crop_by_transform_mat, get_perspective_transform
 
@@ -164,13 +164,17 @@ def apply_transform(
             )
         if self.flags["cropping_mode"] == "slice":  # uses advanced slicing to crop
             B, C, _, _ = input.shape
+            x1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 0].numpy()
+            x2 = params["src"].to(device="cpu", dtype=torch.long)[:, 1, 0].numpy() + 1
+            y1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 1].numpy()
+            y2 = params["src"].to(device="cpu", dtype=torch.long)[:, 3, 1].numpy() + 1
+
+            if self.same_on_batch:
+                return input[..., y1[0]:y2[0], x1[0]:x2[0]]
+
             out = torch.empty(B, C, *self.flags["size"], device=input.device, dtype=input.dtype)
             for i in range(B):
-                x1 = int(params["src"][i, 0, 0])
-                x2 = int(params["src"][i, 1, 0]) + 1
-                y1 = int(params["src"][i, 0, 1])
-                y2 = int(params["src"][i, 3, 1]) + 1
-                out[i] = input[i : i + 1, :, y1:y2, x1:x2]
+                out[i] = input[i : i + 1, :, y1[i]:y2[i], x1[i]:x2[i]]
             return out
         raise NotImplementedError(f"Not supported type: {self.flags['cropping_mode']}.")
 
@@ -234,14 +238,18 @@ def forward(
             input_pad = None
 
         if isinstance(input, (tuple, list)):
+            ori_shape = input[0].shape
             input_temp = _transform_input(input[0])
             input_pad = self.compute_padding(input[0].shape) if input_pad is None else input_pad
             _input = (self.precrop_padding(input_temp, input_pad), input[1])
+            _input = _transform_output_shape(_input, ori_shape) if self.keepdim else _input
         else:
             input = cast(torch.Tensor, input)  # TODO: weird that cast is not working under this context.
+            ori_shape = input.shape
             input_temp = _transform_input(input)
             input_pad = self.compute_padding(input_temp.shape) if input_pad is None else input_pad
             _input = self.precrop_padding(input_temp, input_pad)  # type: ignore
+            _input = _transform_output_shape(_input, ori_shape) if self.keepdim else _input
         out = super().forward(_input, params, return_transform)
 
         # Update the actual input size for inverse
diff --git a/kornia/augmentation/base.py b/kornia/augmentation/base.py
index 124aeeb357..5687f7593b 100644
--- a/kornia/augmentation/base.py
+++ b/kornia/augmentation/base.py
@@ -258,4 +258,5 @@ def forward(  # type: ignore
 
         self._params = params
         output = self.apply_func(in_tensor, in_transform, self._params, return_transform)
+
         return _transform_output_shape(output, ori_shape) if self.keepdim else output


Code review comments:
diff_hunk: @@ -110,14 +110,17 @@ def apply_transform(
         if self.flags["cropping_mode"] == "slice":  # uses advanced slicing to crop
             # TODO: implement as separated function `crop_and_resize_iterative`
             B, C, _, _ = input.shape
-            H, W = self.size
-            out = torch.empty(B, C, H, W, device=input.device, dtype=input.dtype)
+            x1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 0].numpy()
+            x2 = params["src"].to(device="cpu", dtype=torch.long)[:, 1, 0].numpy() + 1
+            y1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 1].numpy()
+            y2 = params["src"].to(device="cpu", dtype=torch.long)[:, 3, 1].numpy() + 1
path: kornia/augmentation/_2d/geometric/center_crop.py
text: @author:
```suggestion
            src = torch.as_tensor(params["src"], device="cpu", dtype=torch.long).numpy()
            x1 = src[:, 0, 0]
            x2 = src[:, 1, 0] + 1
            y1 = src[:, 0, 1]
            y2 = src[:, 3, 1] + 1
```

diff_hunk: @@ -110,14 +110,17 @@ def apply_transform(
         if self.flags["cropping_mode"] == "slice":  # uses advanced slicing to crop
             # TODO: implement as separated function `crop_and_resize_iterative`
             B, C, _, _ = input.shape
-            H, W = self.size
-            out = torch.empty(B, C, H, W, device=input.device, dtype=input.dtype)
+            x1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 0].numpy()
path: kornia/augmentation/_2d/geometric/center_crop.py
text: @user1:
It's not boilerplate to try to cast and send to device four times? Since it's a single tensor, it can be done in a single call

diff_hunk: @@ -164,13 +164,17 @@ def apply_transform(
             )
         if self.flags["cropping_mode"] == "slice":  # uses advanced slicing to crop
             B, C, _, _ = input.shape
+            x1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 0].numpy()
+            x2 = params["src"].to(device="cpu", dtype=torch.long)[:, 1, 0].numpy() + 1
+            y1 = params["src"].to(device="cpu", dtype=torch.long)[:, 0, 1].numpy()
+            y2 = params["src"].to(device="cpu", dtype=torch.long)[:, 3, 1].numpy() + 1
path: kornia/augmentation/_2d/geometric/crop.py
text: @author:
```suggestion
            src = torch.as_tensor(params["src"], device="cpu", dtype=torch.long).numpy()
            x1 = src[:, 0, 0]
            x2 = src[:, 1, 0] + 1
            y1 = src[:, 0, 1]
            y2 = src[:, 3, 1] + 1
```