Adds ImageToTensor module and resize for non-batched images (kornia#978)

* adapt resize for non batch

* implement ImageToTensor module

* handle length 2 and 3

kornia/geometry/transform/affwarp.py

@@ -10,6 +10,7 @@
     warp_affine3d, get_projective_transform
 )
 from kornia.utils import _extract_device_dtype
+from kornia.utils.image import _to_bchw
 
 __all__ = [
     "affine",
@@ -544,7 +545,16 @@ def resize(input: torch.Tensor, size: Union[int, Tuple[int, int]],
     if size == input_size:
         return input
 
-    return torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
+    # TODO: find a proper way to handle this cases in the future
+    input_tmp = _to_bchw(input)
+
+    output = torch.nn.functional.interpolate(
+        input_tmp, size=size, mode=interpolation, align_corners=align_corners)
+
+    if len(input.shape) != len(output.shape):
+        output = output.squeeze()
+
+    return output
 
 
 def rescale(

kornia/utils/__init__.py

@@ -1,6 +1,6 @@
 from .one_hot import one_hot
 from .grid import create_meshgrid, create_meshgrid3d
-from .image import tensor_to_image, image_to_tensor
+from .image import tensor_to_image, image_to_tensor, ImageToTensor
 from .pointcloud_io import save_pointcloud_ply, load_pointcloud_ply
 from .draw import draw_rectangle
 from .helpers import _extract_device_dtype
@@ -17,4 +17,5 @@
     "load_pointcloud_ply",
     "draw_rectangle",
     "_extract_device_dtype",
+    "ImageToTensor",
 ]

kornia/utils/image.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import torch
+import torch.nn as nn
 
 
 def image_to_tensor(image: np.ndarray, keepdim: bool = True) -> torch.Tensor:
@@ -141,3 +142,18 @@ def tensor_to_image(tensor: torch.Tensor) -> np.ndarray:
             "Cannot process tensor with shape {}".format(input_shape))
 
     return image
+
+
+class ImageToTensor(nn.Module):
+    """Converts a numpy image to a PyTorch 4d tensor image.
+
+    Args:
+        keepdim (bool): If ``False`` unsqueeze the input image to match the shape
+            :math:`(B, H, W, C)`. Default: ``True``
+    """
+    def __init__(self, keepdim: bool = False):
+        super().__init__()
+        self.keepdim = keepdim
+
+    def forward(self, x: np.ndarray) -> torch.Tensor:
+        return image_to_tensor(x, keepdim=self.keepdim)

test/geometry/transform/test_affine.py

@@ -13,6 +13,11 @@ def test_smoke(self, device, dtype):
         out = kornia.resize(inp, (3, 4))
         assert_allclose(inp, out, atol=1e-4, rtol=1e-4)
 
+    def test_no_batch(self, device, dtype):
+        inp = torch.rand(3, 3, 4, device=device, dtype=dtype)
+        out = kornia.resize(inp, (2, 5))
+        assert out.shape == (3, 2, 5)
+
     def test_upsize(self, device, dtype):
         inp = torch.rand(1, 3, 3, 4, device=device, dtype=dtype)
         out = kornia.resize(inp, (6, 8))

test/utils/test_image.py

@@ -2,6 +2,8 @@
 import numpy as np
 
 import torch
+from torch.testing import assert_allclose
+
 import kornia as kornia
 import kornia.testing as utils  # test utils
 
@@ -43,6 +45,9 @@ def test_image_to_tensor(input_shape, expected):
     assert tensor.shape == expected
     assert isinstance(tensor, torch.Tensor)
 
+    to_tensor = kornia.utils.ImageToTensor(keepdim=False)
+    assert_allclose(tensor, to_tensor(image))
+
 
 @pytest.mark.parametrize("input_shape, expected",
                          [((4, 4), (1, 4, 4)),