Option to avoid squeezing tensors with batch size = 1 when converting to arrays with tensor_to_image()

🚀 Feature

Adding a boolean input argument keepdim to the tensor_to_image() function to avoid squeezing tensors with batch size 1.

Motivation

While the image_to_tensor() function has a keepdim argument that avoids turning (H, W, C ) images into (1, C, H, W) tensors, there is nothing similar in tensor_to_image() to avoid turning tensors (1, C, H, W) into arrays (H, W, C). Currently, there is no way of keeping the batch dimension when it is 1.

Take a look at this piece of code:

import torch
import kornia

t_img = torch.randn((1, 3, 100, 100))
img = kornia.tensor_to_image(t_img)  # shape 100x100x3

It would be nice to have something like this:

import torch
import kornia

t_img = torch.randn((1, 3, 100, 100))
img = kornia.tensor_to_image(t_img, keepdim=True)  # shape 1x100x100x3

Pitch

When keepdim=True, the tensor_to_image() function returns an array whose first dimension is the batch_size regardless of the batch_size being equal to 1 (as it is currently happening when batch_size > 1).

[arubior]

Sure, I can do it. I’ll set the default value of keepdim to False to maintain the current behaviour if the arg is not specified.

[edgarriba]

fixed in #1168