Merge branch 'main' into proto-perf-feature-improvements

Author: vfdev-5

test/test_prototype_transforms.py

@@ -1379,10 +1379,9 @@ def test__transform(self, mocker):
 
 
 class TestRandomShortestSize:
-    def test__get_params(self, mocker):
+    @pytest.mark.parametrize("min_size,max_size", [([5, 9], 20), ([5, 9], None)])
+    def test__get_params(self, min_size, max_size, mocker):
         spatial_size = (3, 10)
-        min_size = [5, 9]
-        max_size = 20
 
         transform = transforms.RandomShortestSize(min_size=min_size, max_size=max_size)
 
@@ -1395,10 +1394,9 @@ def test__get_params(self, mocker):
         assert isinstance(size, tuple) and len(size) == 2
 
         longer = max(size)
-        assert longer <= max_size
-
         shorter = min(size)
-        if longer == max_size:
+        if max_size is not None:
+            assert longer <= max_size
             assert shorter <= max_size
         else:
             assert shorter in min_size

test/test_prototype_transforms_consistency.py

@@ -12,6 +12,7 @@
 import torch
 from prototype_common_utils import (
     ArgsKwargs,
+    assert_close,
     assert_equal,
     make_bounding_box,
     make_detection_mask,
@@ -40,13 +41,15 @@ def __init__(
         make_images_kwargs=None,
         supports_pil=True,
         removed_params=(),
+        closeness_kwargs=None,
     ):
         self.prototype_cls = prototype_cls
         self.legacy_cls = legacy_cls
         self.args_kwargs = args_kwargs
         self.make_images_kwargs = make_images_kwargs or DEFAULT_MAKE_IMAGES_KWARGS
         self.supports_pil = supports_pil
         self.removed_params = removed_params
+        self.closeness_kwargs = closeness_kwargs or dict(rtol=0, atol=0)
 
 
 # These are here since both the prototype and legacy transform need to be constructed with the same random parameters
@@ -491,10 +494,14 @@ def test_signature_consistency(config):
     assert prototype_kinds == legacy_kinds
 
 
-def check_call_consistency(prototype_transform, legacy_transform, images=None, supports_pil=True):
+def check_call_consistency(
+    prototype_transform, legacy_transform, images=None, supports_pil=True, closeness_kwargs=None
+):
     if images is None:
         images = make_images(**DEFAULT_MAKE_IMAGES_KWARGS)
 
+    closeness_kwargs = closeness_kwargs or dict()
+
     for image in images:
         image_repr = f"[{tuple(image.shape)}, {str(image.dtype).rsplit('.')[-1]}]"
 
@@ -520,10 +527,11 @@ def check_call_consistency(prototype_transform, legacy_transform, images=None, s
                 f"`is_simple_tensor` path in `_transform`."
             ) from exc
 
-        assert_equal(
+        assert_close(
             output_prototype_tensor,
             output_legacy_tensor,
             msg=lambda msg: f"Tensor image consistency check failed with: \n\n{msg}",
+            **closeness_kwargs,
         )
 
         try:
@@ -536,10 +544,11 @@ def check_call_consistency(prototype_transform, legacy_transform, images=None, s
                 f"`features.Image` path in `_transform`."
             ) from exc
 
-        assert_equal(
+        assert_close(
             output_prototype_image,
             output_prototype_tensor,
             msg=lambda msg: f"Output for feature and tensor images is not equal: \n\n{msg}",
+            **closeness_kwargs,
         )
 
         if image.ndim == 3 and supports_pil:
@@ -565,10 +574,11 @@ def check_call_consistency(prototype_transform, legacy_transform, images=None, s
                     f"`PIL.Image.Image` path in `_transform`."
                 ) from exc
 
-            assert_equal(
+            assert_close(
                 output_prototype_pil,
                 output_legacy_pil,
                 msg=lambda msg: f"PIL image consistency check failed with: \n\n{msg}",
+                **closeness_kwargs,
             )
 
 
@@ -606,6 +616,7 @@ def test_call_consistency(config, args_kwargs):
         legacy_transform,
         images=make_images(**config.make_images_kwargs),
         supports_pil=config.supports_pil,
+        closeness_kwargs=config.closeness_kwargs,
     )
 
 

torchvision/prototype/transforms/_geometry.py

@@ -730,7 +730,7 @@ class RandomShortestSize(Transform):
     def __init__(
         self,
         min_size: Union[List[int], Tuple[int], int],
-        max_size: int,
+        max_size: Optional[int] = None,
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
         antialias: Optional[bool] = None,
     ):
@@ -744,7 +744,9 @@ def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         orig_height, orig_width = query_spatial_size(flat_inputs)
 
         min_size = self.min_size[int(torch.randint(len(self.min_size), ()))]
-        r = min(min_size / min(orig_height, orig_width), self.max_size / max(orig_height, orig_width))
+        r = min_size / min(orig_height, orig_width)
+        if self.max_size is not None:
+            r = min(r, self.max_size / max(orig_height, orig_width))
 
         new_width = int(orig_width * r)
         new_height = int(orig_height * r)

torchvision/prototype/transforms/_misc.py

@@ -157,7 +157,10 @@ def __init__(self, dtype: Union[torch.dtype, Dict[Type, torch.dtype]]) -> None:
         self.dtype = dtype
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        return inpt.to(self.dtype[type(inpt)])
+        dtype = self.dtype[type(inpt)]
+        if dtype is None:
+            return inpt
+        return inpt.to(dtype=dtype)
 
 
 class RemoveSmallBoundingBoxes(Transform):

torchvision/prototype/transforms/_transform.py

@@ -5,19 +5,15 @@
 import torch
 from torch import nn
 from torch.utils._pytree import tree_flatten, tree_unflatten
-from torchvision.prototype import features
 from torchvision.prototype.transforms._utils import _isinstance
 from torchvision.utils import _log_api_usage_once
 
 
 class Transform(nn.Module):
 
     # Class attribute defining transformed types. Other types are passed-through without any transformation
-    _transformed_types: Tuple[Union[Type, Callable[[Any], bool]], ...] = (
-        features.is_simple_tensor,
-        features._Feature,
-        PIL.Image.Image,
-    )
+    # We support both Types and callables that are able to do further checks on the type of the input.
+    _transformed_types: Tuple[Union[Type, Callable[[Any], bool]], ...] = (torch.Tensor, PIL.Image.Image)
 
     def __init__(self) -> None:
         super().__init__()

torchvision/prototype/transforms/functional/_color.py

@@ -227,7 +227,7 @@ def equalize_image_tensor(image: torch.Tensor) -> torch.Tensor:
     if image.numel() == 0:
         return image
 
-    return _equalize_image_tensor_vec(image.view(-1, height, width)).view(image.shape)
+    return _equalize_image_tensor_vec(image.view(-1, height, width)).reshape(image.shape)
 
 
 equalize_image_pil = _FP.equalize

torchvision/transforms/functional_tensor.py

@@ -875,7 +875,7 @@ def _scale_channel(img_chan: Tensor) -> Tensor:
     if img_chan.is_cuda:
         hist = torch.histc(img_chan.to(torch.float32), bins=256, min=0, max=255)
     else:
-        hist = torch.bincount(img_chan.view(-1), minlength=256)
+        hist = torch.bincount(img_chan.reshape(-1), minlength=256)
 
     nonzero_hist = hist[hist != 0]
     step = torch.div(nonzero_hist[:-1].sum(), 255, rounding_mode="floor")