[fbsync] Replace get_image_size/num_channels with get_dimensions (#5487)

Summary:
* Replace get_image_size/num_channels with get_image_dims

* Reduce verbosity

* Fix JIT-scriptability

* Refactoring

* More refactoring

* Replace all _FP/_FT direct calls.

* Remove usages of get_image_size and get_image_num_channels from code-base.

* Fix JIT issues

* Adding missing assertion.

Reviewed By: NicolasHug

Differential Revision: D34579514

fbshipit-source-id: 851038b155279541836f2f3228a19f1d0239af57

Author: prabhat00155

docs/source/transforms.rst

@@ -270,6 +270,7 @@ you can use a functional transform to build transform classes with custom behavi
     erase
     five_crop
     gaussian_blur
+    get_dimensions
     get_image_num_channels
     get_image_size
     hflip

references/classification/transforms.py

@@ -141,7 +141,7 @@ def forward(self, batch: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
 
         # Implemented as on cutmix paper, page 12 (with minor corrections on typos).
         lambda_param = float(torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0])
-        W, H = F.get_image_size(batch)
+        _, H, W = F.get_dimensions(batch)
 
         r_x = torch.randint(W, (1,))
         r_y = torch.randint(H, (1,))

references/detection/transforms.py

@@ -34,7 +34,7 @@ def forward(
         if torch.rand(1) < self.p:
             image = F.hflip(image)
             if target is not None:
-                width, _ = F.get_image_size(image)
+                _, _, width = F.get_dimensions(image)
                 target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
                 if "masks" in target:
                     target["masks"] = target["masks"].flip(-1)
@@ -107,7 +107,7 @@ def forward(
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
-        orig_w, orig_h = F.get_image_size(image)
+        _, orig_h, orig_w = F.get_dimensions(image)
 
         while True:
             # sample an option
@@ -192,7 +192,7 @@ def forward(
         if torch.rand(1) >= self.p:
             return image, target
 
-        orig_w, orig_h = F.get_image_size(image)
+        _, orig_h, orig_w = F.get_dimensions(image)
 
         r = self.side_range[0] + torch.rand(1) * (self.side_range[1] - self.side_range[0])
         canvas_width = int(orig_w * r)
@@ -270,7 +270,7 @@ def forward(
                 image = self._contrast(image)
 
         if r[6] < self.p:
-            channels = F.get_image_num_channels(image)
+            channels, _, _ = F.get_dimensions(image)
             permutation = torch.randperm(channels)
 
             is_pil = F._is_pil_image(image)
@@ -317,7 +317,7 @@ def forward(
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
-        orig_width, orig_height = F.get_image_size(image)
+        _, orig_height, orig_width = F.get_dimensions(image)
 
         r = self.scale_range[0] + torch.rand(1) * (self.scale_range[1] - self.scale_range[0])
         new_width = int(self.target_size[1] * r)

test/test_functional_tensor.py

@@ -29,7 +29,7 @@
 
 
 @pytest.mark.parametrize("device", cpu_and_gpu())
-@pytest.mark.parametrize("fn", [F.get_image_size, F.get_image_num_channels])
+@pytest.mark.parametrize("fn", [F.get_image_size, F.get_image_num_channels, F.get_dimensions])
 def test_image_sizes(device, fn):
     script_F = torch.jit.script(fn)
 
@@ -1020,7 +1020,9 @@ def test_resized_crop(device, mode):
 @pytest.mark.parametrize(
     "func, args",
     [
+        (F_t.get_dimensions, ()),
         (F_t.get_image_size, ()),
+        (F_t.get_image_num_channels, ()),
         (F_t.vflip, ()),
         (F_t.hflip, ()),
         (F_t.crop, (1, 2, 4, 5)),

torchvision/prototype/transforms/__init__.py

@@ -8,7 +8,7 @@
 from ._auto_augment import RandAugment, TrivialAugmentWide, AutoAugment
 from ._container import Compose, RandomApply, RandomChoice, RandomOrder
 from ._geometry import HorizontalFlip, Resize, CenterCrop, RandomResizedCrop
-from ._meta_conversion import ConvertBoundingBoxFormat, ConvertImageDtype, ConvertImageColorSpace
+from ._meta import ConvertBoundingBoxFormat, ConvertImageDtype, ConvertImageColorSpace
 from ._misc import Identity, Normalize, ToDtype, Lambda
 from ._presets import CocoEval, ImageNetEval, VocEval, Kinect400Eval, RaftEval
 from ._type_conversion import DecodeImage, LabelToOneHot

torchvision/prototype/transforms/_augment.py

@@ -7,7 +7,7 @@
 from torchvision.prototype import features
 from torchvision.prototype.transforms import Transform, functional as F
 
-from ._utils import query_image
+from ._utils import query_image, get_image_dimensions
 
 
 class RandomErasing(Transform):
@@ -41,8 +41,7 @@ def __init__(
 
     def _get_params(self, sample: Any) -> Dict[str, Any]:
         image = query_image(sample)
-        img_c = F.get_image_num_channels(image)
-        img_w, img_h = F.get_image_size(image)
+        img_c, img_h, img_w = get_image_dimensions(image)
 
         if isinstance(self.value, (int, float)):
             value = [self.value]
@@ -138,7 +137,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         lam = float(self._dist.sample(()))
 
         image = query_image(sample)
-        W, H = F.get_image_size(image)
+        _, H, W = get_image_dimensions(image)
 
         r_x = torch.randint(W, ())
         r_y = torch.randint(H, ())

torchvision/prototype/transforms/_auto_augment.py

@@ -7,7 +7,7 @@
 from torchvision.prototype.transforms import Transform, InterpolationMode, AutoAugmentPolicy, functional as F
 from torchvision.prototype.utils._internal import apply_recursively
 
-from ._utils import query_image
+from ._utils import query_image, get_image_dimensions
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -47,7 +47,7 @@ def dispatch(
                 return input
 
         image = query_image(sample)
-        num_channels = F.get_image_num_channels(image)
+        num_channels, *_ = get_image_dimensions(image)
 
         fill = self.fill
         if isinstance(fill, (int, float)):
@@ -160,8 +160,8 @@ class AutoAugment(_AutoAugmentBase):
     _AUGMENTATION_SPACE = {
         "ShearX": (lambda num_bins, image_size: torch.linspace(0.0, 0.3, num_bins), True),
         "ShearY": (lambda num_bins, image_size: torch.linspace(0.0, 0.3, num_bins), True),
-        "TranslateX": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
-        "TranslateY": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
+        "TranslateX": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
+        "TranslateY": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
         "Rotate": (lambda num_bins, image_size: torch.linspace(0.0, 30.0, num_bins), True),
         "Brightness": (lambda num_bins, image_size: torch.linspace(0.0, 0.9, num_bins), True),
         "Color": (lambda num_bins, image_size: torch.linspace(0.0, 0.9, num_bins), True),
@@ -278,7 +278,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        image_size = F.get_image_size(image)
+        _, height, width = get_image_dimensions(image)
 
         policy = self._policies[int(torch.randint(len(self._policies), ()))]
 
@@ -288,7 +288,7 @@ def forward(self, *inputs: Any) -> Any:
 
             magnitudes_fn, signed = self._AUGMENTATION_SPACE[transform_id]
 
-            magnitudes = magnitudes_fn(10, image_size)
+            magnitudes = magnitudes_fn(10, (height, width))
             if magnitudes is not None:
                 magnitude = float(magnitudes[magnitude_idx])
                 if signed and torch.rand(()) <= 0.5:
@@ -306,8 +306,8 @@ class RandAugment(_AutoAugmentBase):
         "Identity": (lambda num_bins, image_size: None, False),
         "ShearX": (lambda num_bins, image_size: torch.linspace(0.0, 0.3, num_bins), True),
         "ShearY": (lambda num_bins, image_size: torch.linspace(0.0, 0.3, num_bins), True),
-        "TranslateX": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
-        "TranslateY": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
+        "TranslateX": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
+        "TranslateY": (lambda num_bins, image_size: torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
         "Rotate": (lambda num_bins, image_size: torch.linspace(0.0, 30.0, num_bins), True),
         "Brightness": (lambda num_bins, image_size: torch.linspace(0.0, 0.9, num_bins), True),
         "Color": (lambda num_bins, image_size: torch.linspace(0.0, 0.9, num_bins), True),
@@ -334,12 +334,12 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        image_size = F.get_image_size(image)
+        _, height, width = get_image_dimensions(image)
 
         for _ in range(self.num_ops):
             transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE)
 
-            magnitudes = magnitudes_fn(self.num_magnitude_bins, image_size)
+            magnitudes = magnitudes_fn(self.num_magnitude_bins, (height, width))
             if magnitudes is not None:
                 magnitude = float(magnitudes[int(torch.randint(self.num_magnitude_bins, ()))])
                 if signed and torch.rand(()) <= 0.5:
@@ -383,11 +383,11 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        image_size = F.get_image_size(image)
+        _, height, width = get_image_dimensions(image)
 
         transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE)
 
-        magnitudes = magnitudes_fn(self.num_magnitude_bins, image_size)
+        magnitudes = magnitudes_fn(self.num_magnitude_bins, (height, width))
         if magnitudes is not None:
             magnitude = float(magnitudes[int(torch.randint(self.num_magnitude_bins, ()))])
             if signed and torch.rand(()) <= 0.5:

torchvision/prototype/transforms/_geometry.py

@@ -8,7 +8,7 @@
 from torchvision.prototype.transforms import Transform, InterpolationMode, functional as F
 from torchvision.transforms.transforms import _setup_size, _interpolation_modes_from_int
 
-from ._utils import query_image
+from ._utils import query_image, get_image_dimensions
 
 
 class HorizontalFlip(Transform):
@@ -109,7 +109,7 @@ def __init__(
 
     def _get_params(self, sample: Any) -> Dict[str, Any]:
         image = query_image(sample)
-        width, height = F.get_image_size(image)
+        _, height, width = get_image_dimensions(image)
         area = height * width
 
         log_ratio = torch.log(torch.tensor(self.ratio))

torchvision/prototype/transforms/_meta_conversion.py renamed to torchvision/prototype/transforms/_meta.py

@@ -1,10 +1,12 @@
-from typing import Any, Optional, Union
+from typing import Any, Optional, Tuple, Union
 
 import PIL.Image
 import torch
 from torchvision.prototype import features
 from torchvision.prototype.utils._internal import query_recursively
 
+from .functional._meta import get_dimensions_image_tensor, get_dimensions_image_pil
+
 
 def query_image(sample: Any) -> Union[PIL.Image.Image, torch.Tensor, features.Image]:
     def fn(input: Any) -> Optional[Union[PIL.Image.Image, torch.Tensor, features.Image]]:
@@ -17,3 +19,16 @@ def fn(input: Any) -> Optional[Union[PIL.Image.Image, torch.Tensor, features.Ima
         return next(query_recursively(fn, sample))
     except StopIteration:
         raise TypeError("No image was found in the sample")
+
+
+def get_image_dimensions(image: Union[PIL.Image.Image, torch.Tensor, features.Image]) -> Tuple[int, int, int]:
+    if isinstance(image, features.Image):
+        channels = image.num_channels
+        height, width = image.image_size
+    elif isinstance(image, torch.Tensor):
+        channels, height, width = get_dimensions_image_tensor(image)
+    elif isinstance(image, PIL.Image.Image):
+        channels, height, width = get_dimensions_image_pil(image)
+    else:
+        raise TypeError(f"unable to get image dimensions from object of type {type(image).__name__}")
+    return channels, height, width