More refactoring

Author: datumbox

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, get_image_dims
+from ._utils import query_image, get_image_dimensions
 
 
 class RandomErasing(Transform):
@@ -41,7 +41,7 @@ def __init__(
 
     def _get_params(self, sample: Any) -> Dict[str, Any]:
         image = query_image(sample)
-        img_c, img_h, img_w = get_image_dims(image)
+        img_c, img_h, img_w = get_image_dimensions(image)
 
         if isinstance(self.value, (int, float)):
             value = [self.value]
@@ -137,7 +137,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         lam = float(self._dist.sample(()))
 
         image = query_image(sample)
-        _, H, W = get_image_dims(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, get_image_dims
+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, *_ = get_image_dims(image)
+        num_channels, *_ = get_image_dimensions(image)
 
         fill = self.fill
         if isinstance(fill, (int, float)):
@@ -278,7 +278,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        _, height, width = get_image_dims(image)
+        _, height, width = get_image_dimensions(image)
 
         policy = self._policies[int(torch.randint(len(self._policies), ()))]
 
@@ -334,7 +334,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        _, height, width = get_image_dims(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)
@@ -383,7 +383,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
 
         image = query_image(sample)
-        _, height, width = get_image_dims(image)
+        _, height, width = get_image_dimensions(image)
 
         transform_id, (magnitudes_fn, signed) = self._get_random_item(self._AUGMENTATION_SPACE)
 

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, get_image_dims
+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)
-        _, height, width = get_image_dims(image)
+        _, height, width = get_image_dimensions(image)
         area = height * width
 
         log_ratio = torch.log(torch.tensor(self.ratio))

torchvision/prototype/transforms/_utils.py

@@ -20,14 +20,14 @@ def fn(input: Any) -> Optional[Union[PIL.Image.Image, torch.Tensor, features.Ima
         raise TypeError("No image was found in the sample")
 
 
-def get_image_dims(image: Union[PIL.Image.Image, torch.Tensor, features.Image]) -> Tuple[int, int, int]:
+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 = _FT.get_image_dims(image)
+        channels, height, width = _FT.get_dimensions(image)
     elif isinstance(image, PIL.Image.Image):
-        channels, height, width = _FP.get_image_dims(image)
+        channels, height, width = _FP.get_dimensions(image)
     else:
         raise TypeError(f"unable to get image dimensions from object of type {type(image).__name__}")
     return channels, height, width

torchvision/prototype/transforms/functional/_geometry.py

@@ -39,7 +39,7 @@ def resize_image_tensor(
     antialias: Optional[bool] = None,
 ) -> torch.Tensor:
     new_height, new_width = size
-    num_channels, old_height, old_width = _FT.get_image_dims(image)
+    num_channels, old_height, old_width = _FT.get_dimensions(image)
     batch_shape = image.shape[:-3]
     return _FT.resize(
         image.reshape((-1, num_channels, old_height, old_width)),
@@ -141,7 +141,7 @@ def affine_image_tensor(
 
     center_f = [0.0, 0.0]
     if center is not None:
-        _, height, width = _FT.get_image_dims(img)
+        _, height, width = _FT.get_dimensions(img)
         # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
         center_f = [1.0 * (c - s * 0.5) for c, s in zip(center, (width, height))]
 
@@ -167,7 +167,7 @@ def affine_image_pil(
     # it is visually better to estimate the center without 0.5 offset
     # otherwise image rotated by 90 degrees is shifted vs output image of torch.rot90 or F_t.affine
     if center is None:
-        _, height, width = _FP.get_image_dims(img)
+        _, height, width = _FP.get_dimensions(img)
         center = [width * 0.5, height * 0.5]
     matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear)
 
@@ -184,7 +184,7 @@ def rotate_image_tensor(
 ) -> torch.Tensor:
     center_f = [0.0, 0.0]
     if center is not None:
-        _, height, width = _FT.get_image_dims(img)
+        _, height, width = _FT.get_dimensions(img)
         # Center values should be in pixel coordinates but translated such that (0, 0) corresponds to image center.
         center_f = [1.0 * (c - s * 0.5) for c, s in zip(center, (width, height))]
 
@@ -260,13 +260,13 @@ def _center_crop_compute_crop_anchor(
 
 def center_crop_image_tensor(img: torch.Tensor, output_size: List[int]) -> torch.Tensor:
     crop_height, crop_width = _center_crop_parse_output_size(output_size)
-    _, image_height, image_width = _FT.get_image_dims(img)
+    _, image_height, image_width = _FT.get_dimensions(img)
 
     if crop_height > image_height or crop_width > image_width:
         padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width)
         img = pad_image_tensor(img, padding_ltrb, fill=0)
 
-        _, image_height, image_width = _FT.get_image_dims(img)
+        _, image_height, image_width = _FT.get_dimensions(img)
         if crop_width == image_width and crop_height == image_height:
             return img
 
@@ -276,13 +276,13 @@ def center_crop_image_tensor(img: torch.Tensor, output_size: List[int]) -> torch
 
 def center_crop_image_pil(img: PIL.Image.Image, output_size: List[int]) -> PIL.Image.Image:
     crop_height, crop_width = _center_crop_parse_output_size(output_size)
-    _, image_height, image_width = _FP.get_image_dims(img)
+    _, image_height, image_width = _FP.get_dimensions(img)
 
     if crop_height > image_height or crop_width > image_width:
         padding_ltrb = _center_crop_compute_padding(crop_height, crop_width, image_height, image_width)
         img = pad_image_pil(img, padding_ltrb, fill=0)
 
-        _, image_height, image_width = _FP.get_image_dims(img)
+        _, image_height, image_width = _FP.get_dimensions(img)
         if crop_width == image_width and crop_height == image_height:
             return img
 

torchvision/prototype/transforms/functional/_misc.py

@@ -2,10 +2,13 @@
 
 import PIL.Image
 import torch
-from torchvision.transforms import functional_tensor as _FT
+from torchvision.transforms import functional_tensor as _FT, functional_pil as _FP
 from torchvision.transforms.functional import to_tensor, to_pil_image
 
 
+get_dimensions_image_tensor = _FT.get_dimensions
+get_dimensions_image_pil = _FP.get_dimensions
+
 normalize_image_tensor = _FT.normalize
 
 

torchvision/transforms/functional_pil.py

@@ -21,7 +21,7 @@ def _is_pil_image(img: Any) -> bool:
 
 
 @torch.jit.unused
-def get_image_dims(img: Any) -> List[int]:
+def get_dimensions(img: Any) -> List[int]:
     if _is_pil_image(img):
         channels = len(img.getbands())
         width, height = img.size

torchvision/transforms/functional_tensor.py

@@ -21,7 +21,7 @@ def _assert_threshold(img: Tensor, threshold: float) -> None:
         raise TypeError("Threshold should be less than bound of img.")
 
 
-def get_image_dims(img: Tensor) -> List[int]:
+def get_dimensions(img: Tensor) -> List[int]:
     _assert_image_tensor(img)
     channels = 1 if img.ndim == 2 else img.shape[-3]
     height, width = img.shape[-2:]