Repository: pytorch/vision

Issue:
Simplify dimension checks on functional_tensor.py
## 🚀 Feature
The functional_tensor.py file makes use of the private method `_is_tensor_a_torch_image` in every public operator to check its dimensions:
https://github.com/pytorch/vision/blob/dab475720f116c42fc80e437b25496dfc94d2a8a/torchvision/transforms/functional_tensor.py#L10-L11

Examples:
https://github.com/pytorch/vision/blob/dab475720f116c42fc80e437b25496dfc94d2a8a/torchvision/transforms/functional_tensor.py#L146-L147
https://github.com/pytorch/vision/blob/dab475720f116c42fc80e437b25496dfc94d2a8a/torchvision/transforms/functional_tensor.py#L166-L167

This check is repetitive and reduces the code readability. We should fix this by using ~decorators. See https://github.com/pytorch/vision/pull/3123#discussion_r540840674 for details.~ assertions. 


cc @vfdev-5


Discussion:
FYI I'm not sure decorators are supported in torchscript
> FYI I'm not sure decorators are supported in torchscript

The following works
```python
import torch
print(torch.__version__)
import typing


def check_tensor(func: typing.Callable) -> typing.Callable:
    
    def wrapper(t: torch.Tensor, p: int = 5):
        if not isinstance(t, torch.Tensor):
            raise TypeError()
        print("Inside decorator")
        return func(t, p)
    
    return wrapper


@check_tensor
def foo(t: torch.Tensor, p: int = 5) -> torch.Tensor:
    for i in range(10):
        if i < p:
            continue
        t += i
    return t


print(foo(torch.tensor([0.0])))
sfoo = torch.jit.script(foo)
print(sfoo(torch.tensor([0.0])))

> 1.8.0.dev20201211
> Inside decorator
> tensor([35.])
> Inside decorator
> tensor([35.])
```
Great!
Somebody working on this feature?
@avijit9 Not at the moment. If you are interested in sending a PR that would be awesome!
How to deal with multiple arguments? For example, the crop function takes multiple arguments. But if we use a decorator with `*args, **kwargs` I'm getting the following error - 

`FAILED test/test_decorator.py::Tester::test_crop - torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:`

The decorator I implemented is shown below - 

```
def check_tensor_torch_image(f: typing.Callable) -> typing.Callable:
    def wrapper(*args, **kwargs):
        if args[0].ndim < 2:
            raise TypeError("Tensor is not a torch image.")
        return f(*args, **kwargs)
    return wrapper
```

The unit test is as follows - 


```
	def test_crop(self):
		scripted_fn = torch.jit.script(F_t.crop)
		shape = (10,)
		tensor = torch.rand(*shape, dtype=torch.float, device=self.device)
		with self.assertRaises(Exception) as context:
			scripted_fn(tensor)
		self.assertTrue('Tensor is not a torch image.' in str(context.exception))
```

This is a known issue - https://github.com/pytorch/pytorch/issues/29637

cc - @datumbox 
Minimal code to reproduce this -
```python

import torch
print(torch.__version__)
import typing


def check_tensor(func: typing.Callable) -> typing.Callable:
    
    def wrapper(*args, **kwargs):
        if not isinstance(args[0], torch.Tensor):
            raise TypeError()
        print("Inside decorator")
        return func(*args, **kwargs)
    
    return wrapper


@check_tensor
def foo(t: torch.Tensor, p: int = 5) -> torch.Tensor:
    for i in range(10):
        if i < p:
            continue
        t += i
    return t


> print(foo(torch.tensor([0.0])))
> sfoo = torch.jit.script(foo)
> print(sfoo(torch.tensor([0.0])))

```
@avijit9 Thank you for the detailed analysis!

@vfdev-5 Do you have any suggestion to work around the limitation? I suppose it's possible to implement a couple of versions of the decorator since most of the methods in functional_tensor.py have limited number of arguments but I think that would lead to more messy code which defeats the purpose of your original recommendation. Thoughts? 
@datumbox yes, seems like it is a blocker for decorator usage. On the other hand, the goal of this issue is to refactor a bit the current code base by removing explicit type checking like
```python
 if not _is_tensor_a_torch_image(img): 
     raise TypeError('tensor is not a torch image.') 
```
Using decorator we remove 2 lines and add 1. I propose to just create a simple helper methods and reuse them with the similar number of lines gain:
```python
def _assert_image_tensor(img):
    if not _is_tensor_a_torch_image(img):
        raise TypeError("tensor is not a torch image.")    

def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor:
    """PRIVATE METHOD. Crop the given Image Tensor.
    """
    _assert_image_tensor(img)

    return img[..., top:top + height, left:left + width]

```

What do you think ?
@vfdev-5 agreed, that's still an improvement.

@avijit9 would you be still interested to send a PR based on the above proposal?
@datumbox  Sure! I'll send a PR soon.

Resolving patch:
diff --git a/torchvision/transforms/functional_tensor.py b/torchvision/transforms/functional_tensor.py
--- a/torchvision/transforms/functional_tensor.py
+++ b/torchvision/transforms/functional_tensor.py
@@ -11,11 +11,15 @@ def _is_tensor_a_torch_image(x: Tensor) -> bool:
     return x.ndim >= 2
 
 
+def _assert_image_tensor(img):
+    if not _is_tensor_a_torch_image(img):
+        raise TypeError("Tensor is not a torch image.")
+
+
 def _get_image_size(img: Tensor) -> List[int]:
     """Returns (w, h) of tensor image"""
-    if _is_tensor_a_torch_image(img):
-        return [img.shape[-1], img.shape[-2]]
-    raise TypeError("Unexpected input type")
+    _assert_image_tensor(img)
+    return [img.shape[-1], img.shape[-2]]
 
 
 def _get_image_num_channels(img: Tensor) -> int:
@@ -143,8 +147,7 @@ def vflip(img: Tensor) -> Tensor:
     Returns:
         Tensor:  Vertically flipped image Tensor.
     """
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     return img.flip(-2)
 
@@ -163,8 +166,7 @@ def hflip(img: Tensor) -> Tensor:
     Returns:
         Tensor:  Horizontally flipped image Tensor.
     """
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     return img.flip(-1)
 
@@ -187,8 +189,7 @@ def crop(img: Tensor, top: int, left: int, height: int, width: int) -> Tensor:
     Returns:
         Tensor: Cropped image.
     """
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError("tensor is not a torch image.")
+    _assert_image_tensor(img)
 
     return img[..., top:top + height, left:left + width]
 
@@ -254,8 +255,7 @@ def adjust_brightness(img: Tensor, brightness_factor: float) -> Tensor:
     if brightness_factor < 0:
         raise ValueError('brightness_factor ({}) is not non-negative.'.format(brightness_factor))
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     _assert_channels(img, [1, 3])
 
@@ -282,8 +282,7 @@ def adjust_contrast(img: Tensor, contrast_factor: float) -> Tensor:
     if contrast_factor < 0:
         raise ValueError('contrast_factor ({}) is not non-negative.'.format(contrast_factor))
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     _assert_channels(img, [3])
 
@@ -326,9 +325,11 @@ def adjust_hue(img: Tensor, hue_factor: float) -> Tensor:
     if not (-0.5 <= hue_factor <= 0.5):
         raise ValueError('hue_factor ({}) is not in [-0.5, 0.5].'.format(hue_factor))
 
-    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
+    if not (isinstance(img, torch.Tensor)):
         raise TypeError('Input img should be Tensor image')
 
+    _assert_image_tensor(img)
+
     _assert_channels(img, [3])
 
     orig_dtype = img.dtype
@@ -367,8 +368,7 @@ def adjust_saturation(img: Tensor, saturation_factor: float) -> Tensor:
     if saturation_factor < 0:
         raise ValueError('saturation_factor ({}) is not non-negative.'.format(saturation_factor))
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     _assert_channels(img, [3])
 
@@ -447,8 +447,7 @@ def center_crop(img: Tensor, output_size: BroadcastingList2[int]) -> Tensor:
         "Please, use ``F.center_crop`` instead."
     )
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     _, image_width, image_height = img.size()
     crop_height, crop_width = output_size
@@ -497,8 +496,7 @@ def five_crop(img: Tensor, size: BroadcastingList2[int]) -> List[Tensor]:
         "Please, use ``F.five_crop`` instead."
     )
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
 
@@ -553,8 +551,7 @@ def ten_crop(img: Tensor, size: BroadcastingList2[int], vertical_flip: bool = Fa
         "Please, use ``F.ten_crop`` instead."
     )
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
     first_five = five_crop(img, size)
@@ -703,8 +700,7 @@ def pad(img: Tensor, padding: List[int], fill: int = 0, padding_mode: str = "con
     Returns:
         Tensor: Padded image.
     """
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError("tensor is not a torch image.")
+    _assert_image_tensor(img)
 
     if not isinstance(padding, (int, tuple, list)):
         raise TypeError("Got inappropriate padding arg")
@@ -796,8 +792,7 @@ def resize(img: Tensor, size: List[int], interpolation: str = "bilinear") -> Ten
     Returns:
         Tensor: Resized image.
     """
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError("tensor is not a torch image.")
+    _assert_image_tensor(img)
 
     if not isinstance(size, (int, tuple, list)):
         raise TypeError("Got inappropriate size arg")
@@ -855,8 +850,11 @@ def _assert_grid_transform_inputs(
         supported_interpolation_modes: List[str],
         coeffs: Optional[List[float]] = None,
 ):
-    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
-        raise TypeError("Input img should be Tensor Image")
+
+    if not (isinstance(img, torch.Tensor)):
+        raise TypeError("Input img should be Tensor")
+
+    _assert_image_tensor(img)
 
     if matrix is not None and not isinstance(matrix, list):
         raise TypeError("Argument matrix should be a list")
@@ -1112,8 +1110,11 @@ def perspective(
     Returns:
         Tensor: transformed image.
     """
-    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
-        raise TypeError('Input img should be Tensor Image')
+
+    if not (isinstance(img, torch.Tensor)):
+        raise TypeError('Input img should be Tensor.')
+
+    _assert_image_tensor(img)
 
     _assert_grid_transform_inputs(
         img,
@@ -1165,8 +1166,11 @@ def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: List[float]) -> Te
     Returns:
         Tensor: An image that is blurred using gaussian kernel of given parameters
     """
-    if not (isinstance(img, torch.Tensor) or _is_tensor_a_torch_image(img)):
-        raise TypeError('img should be Tensor Image. Got {}'.format(type(img)))
+
+    if not (isinstance(img, torch.Tensor)):
+        raise TypeError('img should be Tensor. Got {}'.format(type(img)))
+
+    _assert_image_tensor(img)
 
     dtype = img.dtype if torch.is_floating_point(img) else torch.float32
     kernel = _get_gaussian_kernel2d(kernel_size, sigma, dtype=dtype, device=img.device)
@@ -1184,8 +1188,8 @@ def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: List[float]) -> Te
 
 
 def invert(img: Tensor) -> Tensor:
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+
+    _assert_image_tensor(img)
 
     if img.ndim < 3:
         raise TypeError("Input image tensor should have at least 3 dimensions, but found {}".format(img.ndim))
@@ -1197,8 +1201,8 @@ def invert(img: Tensor) -> Tensor:
 
 
 def posterize(img: Tensor, bits: int) -> Tensor:
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+
+    _assert_image_tensor(img)
 
     if img.ndim < 3:
         raise TypeError("Input image tensor should have at least 3 dimensions, but found {}".format(img.ndim))
@@ -1211,8 +1215,8 @@ def posterize(img: Tensor, bits: int) -> Tensor:
 
 
 def solarize(img: Tensor, threshold: float) -> Tensor:
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+
+    _assert_image_tensor(img)
 
     if img.ndim < 3:
         raise TypeError("Input image tensor should have at least 3 dimensions, but found {}".format(img.ndim))
@@ -1245,8 +1249,7 @@ def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor:
     if sharpness_factor < 0:
         raise ValueError('sharpness_factor ({}) is not non-negative.'.format(sharpness_factor))
 
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+    _assert_image_tensor(img)
 
     _assert_channels(img, [1, 3])
 
@@ -1257,8 +1260,8 @@ def adjust_sharpness(img: Tensor, sharpness_factor: float) -> Tensor:
 
 
 def autocontrast(img: Tensor) -> Tensor:
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+
+    _assert_image_tensor(img)
 
     if img.ndim < 3:
         raise TypeError("Input image tensor should have at least 3 dimensions, but found {}".format(img.ndim))
@@ -1297,8 +1300,8 @@ def _equalize_single_image(img: Tensor) -> Tensor:
 
 
 def equalize(img: Tensor) -> Tensor:
-    if not _is_tensor_a_torch_image(img):
-        raise TypeError('tensor is not a torch image.')
+
+    _assert_image_tensor(img)
 
     if not (3 <= img.ndim <= 4):
         raise TypeError("Input image tensor should have 3 or 4 dimensions, but found {}".format(img.ndim))


Test patch:
diff --git a/test/test_functional_tensor.py b/test/test_functional_tensor.py
--- a/test/test_functional_tensor.py
+++ b/test/test_functional_tensor.py
@@ -13,6 +13,8 @@
 
 from common_utils import TransformsTester
 
+from typing import Dict, List, Tuple
+
 
 NEAREST, BILINEAR, BICUBIC = InterpolationMode.NEAREST, InterpolationMode.BILINEAR, InterpolationMode.BICUBIC
 
@@ -34,6 +36,28 @@ def _test_fn_on_batch(self, batch_tensors, fn, **fn_kwargs):
         s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs)
         self.assertTrue(transformed_batch.allclose(s_transformed_batch))
 
+    def test_assert_image_tensor(self):
+        shape = (100,)
+        tensor = torch.rand(*shape, dtype=torch.float, device=self.device)
+
+        list_of_methods = [(F_t._get_image_size, (tensor, )), (F_t.vflip, (tensor, )),
+                           (F_t.hflip, (tensor, )), (F_t.crop, (tensor, 1, 2, 4, 5)),
+                           (F_t.adjust_brightness, (tensor, 0.)), (F_t.adjust_contrast, (tensor, 1.)),
+                           (F_t.adjust_hue, (tensor, -0.5)), (F_t.adjust_saturation, (tensor, 2.)),
+                           (F_t.center_crop, (tensor, [10, 11])), (F_t.five_crop, (tensor, [10, 11])),
+                           (F_t.ten_crop, (tensor, [10, 11])), (F_t.pad, (tensor, [2, ], 2, "constant")),
+                           (F_t.resize, (tensor, [10, 11])), (F_t.perspective, (tensor, [0.2, ])),
+                           (F_t.gaussian_blur, (tensor, (2, 2), (0.7, 0.5))),
+                           (F_t.invert, (tensor, )), (F_t.posterize, (tensor, 0)),
+                           (F_t.solarize, (tensor, 0.3)), (F_t.adjust_sharpness, (tensor, 0.3)),
+                           (F_t.autocontrast, (tensor, )), (F_t.equalize, (tensor, ))]
+
+        for func, args in list_of_methods:
+            with self.assertRaises(Exception) as context:
+                func(*args)
+
+            self.assertTrue('Tensor is not a torch image.' in str(context.exception))
+
     def test_vflip(self):
         script_vflip = torch.jit.script(F.vflip)