Refactoring typing definitions.

Author: datumbox

torchvision/prototype/features/_bounding_box.py

@@ -6,8 +6,7 @@
 from torchvision._utils import StrEnum
 from torchvision.transforms import InterpolationMode  # TODO: this needs to be moved out of transforms
 
-from ._feature import _Feature
-from ._utils import FillType
+from ._feature import _Feature, FillTypeJIT
 
 
 class BoundingBoxFormat(StrEnum):
@@ -116,7 +115,7 @@ def resized_crop(
     def pad(
         self,
         padding: Union[int, Sequence[int]],
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         padding_mode: str = "constant",
     ) -> BoundingBox:
         # This cast does Sequence[int] -> List[int] and is required to make mypy happy
@@ -138,7 +137,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> BoundingBox:
         output = self._F.rotate_bounding_box(
@@ -166,7 +165,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> BoundingBox:
         output = self._F.affine_bounding_box(
@@ -185,7 +184,7 @@ def perspective(
         self,
         perspective_coeffs: List[float],
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> BoundingBox:
         output = self._F.perspective_bounding_box(self, self.format, perspective_coeffs)
         return BoundingBox.new_like(self, output, dtype=output.dtype)
@@ -194,7 +193,7 @@ def elastic(
         self,
         displacement: torch.Tensor,
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> BoundingBox:
         output = self._F.elastic_bounding_box(self, self.format, displacement)
         return BoundingBox.new_like(self, output, dtype=output.dtype)

torchvision/prototype/features/_feature.py

@@ -3,14 +3,14 @@
 from types import ModuleType
 from typing import Any, Callable, cast, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
 
+import PIL.Image
 import torch
 from torch._C import _TensorBase, DisableTorchFunction
 from torchvision.transforms import InterpolationMode
 
-from ._utils import FillType
-
-
 F = TypeVar("F", bound="_Feature")
+FillType = Union[int, float, Sequence[int], Sequence[float], None]
+FillTypeJIT = Union[int, float, List[float], None]
 
 
 def is_simple_tensor(inpt: Any) -> bool:
@@ -152,7 +152,7 @@ def resized_crop(
     def pad(
         self,
         padding: Union[int, List[int]],
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         padding_mode: str = "constant",
     ) -> _Feature:
         return self
@@ -162,7 +162,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> _Feature:
         return self
@@ -174,7 +174,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> _Feature:
         return self
@@ -183,15 +183,15 @@ def perspective(
         self,
         perspective_coeffs: List[float],
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> _Feature:
         return self
 
     def elastic(
         self,
         displacement: torch.Tensor,
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> _Feature:
         return self
 
@@ -230,3 +230,7 @@ def invert(self) -> _Feature:
 
     def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> _Feature:
         return self
+
+
+InputType = Union[torch.Tensor, PIL.Image.Image, _Feature]
+InputTypeJIT = torch.Tensor

torchvision/prototype/features/_image.py

@@ -3,14 +3,14 @@
 import warnings
 from typing import Any, cast, List, Optional, Tuple, Union
 
+import PIL.Image
 import torch
 from torchvision._utils import StrEnum
 from torchvision.transforms.functional import InterpolationMode, to_pil_image
 from torchvision.utils import draw_bounding_boxes, make_grid
 
 from ._bounding_box import BoundingBox
-from ._feature import _Feature
-from ._utils import FillType
+from ._feature import _Feature, FillTypeJIT
 
 
 class ColorSpace(StrEnum):
@@ -177,7 +177,7 @@ def resized_crop(
     def pad(
         self,
         padding: Union[int, List[int]],
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         padding_mode: str = "constant",
     ) -> Image:
         output = self._F.pad_image_tensor(self, padding, fill=fill, padding_mode=padding_mode)
@@ -188,7 +188,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> Image:
         output = self._F._geometry.rotate_image_tensor(
@@ -203,7 +203,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> Image:
         output = self._F._geometry.affine_image_tensor(
@@ -222,7 +222,7 @@ def perspective(
         self,
         perspective_coeffs: List[float],
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> Image:
         output = self._F._geometry.perspective_image_tensor(
             self, perspective_coeffs, interpolation=interpolation, fill=fill
@@ -233,7 +233,7 @@ def elastic(
         self,
         displacement: torch.Tensor,
         interpolation: InterpolationMode = InterpolationMode.BILINEAR,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> Image:
         output = self._F._geometry.elastic_image_tensor(self, displacement, interpolation=interpolation, fill=fill)
         return Image.new_like(self, output)
@@ -285,3 +285,11 @@ def invert(self) -> Image:
     def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Image:
         output = self._F.gaussian_blur_image_tensor(self, kernel_size=kernel_size, sigma=sigma)
         return Image.new_like(self, output)
+
+
+ImageType = Union[torch.Tensor, PIL.Image.Image, Image]
+ImageTypeJIT = torch.Tensor
+LegacyImageType = Union[torch.Tensor, PIL.Image.Image]
+LegacyImageTypeJIT = torch.Tensor
+TensorImageType = Union[torch.Tensor, Image]
+TensorImageTypeJIT = torch.Tensor

torchvision/prototype/features/_mask.py

@@ -5,8 +5,7 @@
 import torch
 from torchvision.transforms import InterpolationMode
 
-from ._feature import _Feature
-from ._utils import FillType
+from ._feature import _Feature, FillTypeJIT
 
 
 class Mask(_Feature):
@@ -52,7 +51,7 @@ def resized_crop(
     def pad(
         self,
         padding: Union[int, List[int]],
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         padding_mode: str = "constant",
     ) -> Mask:
         output = self._F.pad_mask(self, padding, padding_mode=padding_mode, fill=fill)
@@ -63,7 +62,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> Mask:
         output = self._F.rotate_mask(self, angle, expand=expand, center=center, fill=fill)
@@ -76,7 +75,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
         center: Optional[List[float]] = None,
     ) -> Mask:
         output = self._F.affine_mask(
@@ -94,7 +93,7 @@ def perspective(
         self,
         perspective_coeffs: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> Mask:
         output = self._F.perspective_mask(self, perspective_coeffs, fill=fill)
         return Mask.new_like(self, output)
@@ -103,7 +102,7 @@ def elastic(
         self,
         displacement: torch.Tensor,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: FillType = None,
+        fill: FillTypeJIT = None,
     ) -> Mask:
         output = self._F.elastic_mask(self, displacement, fill=fill)
         return Mask.new_like(self, output, dtype=output.dtype)

torchvision/prototype/features/_utils.py

@@ -10,8 +10,10 @@
 from torchvision.prototype import features
 from torchvision.prototype.transforms import functional as F, InterpolationMode
 
+from ..features._image import ImageType, TensorImageType
+
 from ._transform import _RandomApplyTransform
-from ._utils import has_any, ImageType, query_chw, TensorImageType
+from ._utils import has_any, query_chw
 
 
 class RandomErasing(_RandomApplyTransform):

torchvision/prototype/transforms/_augment.py

@@ -9,7 +9,10 @@
 from torchvision.prototype.transforms import AutoAugmentPolicy, functional as F, InterpolationMode, Transform
 from torchvision.prototype.transforms.functional._meta import get_chw
 
-from ._utils import _isinstance, _setup_fill_arg, FillType, ImageType
+from ..features._feature import FillType
+from ..features._image import ImageType
+
+from ._utils import _isinstance, _setup_fill_arg
 
 K = TypeVar("K")
 V = TypeVar("V")

torchvision/prototype/transforms/_auto_augment.py

@@ -6,8 +6,10 @@
 from torchvision.prototype import features
 from torchvision.prototype.transforms import functional as F, Transform
 
+from ..features._image import ImageType
+
 from ._transform import _RandomApplyTransform
-from ._utils import ImageType, query_chw
+from ._utils import query_chw
 
 
 class ColorJitter(Transform):

torchvision/prototype/transforms/_color.py

@@ -10,8 +10,10 @@
 from torchvision.transforms import functional as _F
 from typing_extensions import Literal
 
+from ..features._image import ImageType
+
 from ._transform import _RandomApplyTransform
-from ._utils import ImageType, query_chw
+from ._utils import query_chw
 
 
 class ToTensor(Transform):

torchvision/prototype/transforms/_deprecated.py

@@ -12,6 +12,9 @@
 
 from typing_extensions import Literal
 
+from ..features._feature import FillType
+from ..features._image import ImageType
+
 from ._transform import _RandomApplyTransform
 from ._utils import (
     _check_padding_arg,
@@ -20,10 +23,8 @@
     _setup_angle,
     _setup_fill_arg,
     _setup_size,
-    FillType,
     has_all,
     has_any,
-    ImageType,
     query_bounding_box,
     query_chw,
 )