move simple_tensor to features module

Author: pmeier

torchvision/prototype/features/__init__.py

@@ -1,6 +1,6 @@
 from ._bounding_box import BoundingBox, BoundingBoxFormat
 from ._encoded import EncodedData, EncodedImage, EncodedVideo
-from ._feature import _Feature
+from ._feature import _Feature, is_simple_tensor
 from ._image import ColorSpace, Image
 from ._label import Label, OneHotLabel
 from ._segmentation_mask import SegmentationMask

torchvision/prototype/features/_feature.py

@@ -10,6 +10,10 @@
 F = TypeVar("F", bound="_Feature")
 
 
+def is_simple_tensor(inpt: Any) -> bool:
+    return isinstance(inpt, torch.Tensor) and not isinstance(inpt, _Feature)
+
+
 class _Feature(torch.Tensor):
     __F: Optional[ModuleType] = None
 

torchvision/prototype/transforms/_augment.py

@@ -13,7 +13,7 @@
 from torchvision.transforms.functional import InterpolationMode, pil_to_tensor
 
 from ._transform import _RandomApplyTransform
-from ._utils import has_any, is_simple_tensor, query_chw
+from ._utils import has_any, query_chw
 
 
 class RandomErasing(_RandomApplyTransform):
@@ -103,7 +103,7 @@ def __init__(self, *, alpha: float, p: float = 0.5) -> None:
 
     def forward(self, *inpts: Any) -> Any:
         sample = inpts if len(inpts) > 1 else inpts[0]
-        if not (has_any(sample, features.Image, is_simple_tensor) and has_any(sample, features.OneHotLabel)):
+        if not (has_any(sample, features.Image, features.is_simple_tensor) and has_any(sample, features.OneHotLabel)):
             raise TypeError(f"{type(self).__name__}() is only defined for tensor images and one-hot labels.")
         if has_any(sample, features.BoundingBox, features.SegmentationMask, features.Label):
             raise TypeError(
@@ -125,7 +125,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         lam = params["lam"]
-        if isinstance(inpt, features.Image) or is_simple_tensor(inpt):
+        if isinstance(inpt, features.Image) or features.is_simple_tensor(inpt):
             if inpt.ndim < 4:
                 raise ValueError("Need a batch of images")
             output = inpt.clone()
@@ -165,7 +165,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         return dict(box=box, lam_adjusted=lam_adjusted)
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        if isinstance(inpt, features.Image) or is_simple_tensor(inpt):
+        if isinstance(inpt, features.Image) or features.is_simple_tensor(inpt):
             box = params["box"]
             if inpt.ndim < 4:
                 raise ValueError("Need a batch of images")
@@ -275,7 +275,7 @@ def _extract_image_targets(self, flat_sample: List[Any]) -> Tuple[List[Any], Lis
         # with List[image], List[BoundingBox], List[SegmentationMask], List[Label]
         images, bboxes, masks, labels = [], [], [], []
         for obj in flat_sample:
-            if isinstance(obj, features.Image) or is_simple_tensor(obj):
+            if isinstance(obj, features.Image) or features.is_simple_tensor(obj):
                 images.append(obj)
             elif isinstance(obj, PIL.Image.Image):
                 images.append(pil_to_tensor(obj))
@@ -309,7 +309,7 @@ def _insert_outputs(
             elif isinstance(obj, PIL.Image.Image):
                 flat_sample[i] = F.to_image_pil(output_images[c0])
                 c0 += 1
-            elif is_simple_tensor(obj):
+            elif features.is_simple_tensor(obj):
                 flat_sample[i] = output_images[c0]
                 c0 += 1
             elif isinstance(obj, features.BoundingBox):

torchvision/prototype/transforms/_auto_augment.py

@@ -10,7 +10,7 @@
 from torchvision.transforms.autoaugment import AutoAugmentPolicy
 from torchvision.transforms.functional import InterpolationMode, pil_to_tensor, to_pil_image
 
-from ._utils import get_chw, is_simple_tensor
+from ._utils import get_chw
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -46,7 +46,7 @@ def _extract_image(
         sample_flat, _ = tree_flatten(sample)
         images = []
         for id, inpt in enumerate(sample_flat):
-            if isinstance(inpt, (features.Image, PIL.Image.Image)) or is_simple_tensor(inpt):
+            if isinstance(inpt, (features.Image, PIL.Image.Image)) or features.is_simple_tensor(inpt):
                 images.append((id, inpt))
             elif isinstance(inpt, unsupported_types):
                 raise TypeError(f"Inputs of type {type(inpt).__name__} are not supported by {type(self).__name__}()")

torchvision/prototype/transforms/_color.py

@@ -8,7 +8,7 @@
 from torchvision.transforms import functional as _F
 
 from ._transform import _RandomApplyTransform
-from ._utils import is_simple_tensor, query_chw
+from ._utils import query_chw
 
 T = TypeVar("T", features.Image, torch.Tensor, PIL.Image.Image)
 
@@ -112,7 +112,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         )
 
     def _permute_channels(self, inpt: Any, *, permutation: torch.Tensor) -> Any:
-        if not (isinstance(inpt, (features.Image, PIL.Image.Image)) or is_simple_tensor(inpt)):
+        if not (isinstance(inpt, (features.Image, PIL.Image.Image)) or features.is_simple_tensor(inpt)):
             return inpt
 
         image = inpt

torchvision/prototype/transforms/_deprecated.py

@@ -10,7 +10,7 @@
 from typing_extensions import Literal
 
 from ._transform import _RandomApplyTransform
-from ._utils import is_simple_tensor, query_chw
+from ._utils import query_chw
 
 
 class ToTensor(Transform):
@@ -61,7 +61,7 @@ def __init__(self, mode: Optional[str] = None) -> None:
         self.mode = mode
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        if is_simple_tensor(inpt) or isinstance(inpt, (features.Image, np.ndarray)):
+        if features.is_simple_tensor(inpt) or isinstance(inpt, (features.Image, np.ndarray)):
             return _F.to_pil_image(inpt, mode=self.mode)
         else:
             return inpt

torchvision/prototype/transforms/_geometry.py

@@ -15,7 +15,7 @@
 from typing_extensions import Literal
 
 from ._transform import _RandomApplyTransform
-from ._utils import has_all, has_any, is_simple_tensor, query_bounding_box, query_chw
+from ._utils import has_all, has_any, query_bounding_box, query_chw
 
 
 class RandomHorizontalFlip(_RandomApplyTransform):
@@ -700,7 +700,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
         if not (
             has_all(sample, features.BoundingBox)
-            and has_any(sample, PIL.Image.Image, features.Image, is_simple_tensor)
+            and has_any(sample, PIL.Image.Image, features.Image, features.is_simple_tensor)
             and has_any(sample, features.Label, features.OneHotLabel)
         ):
             raise TypeError(
@@ -848,7 +848,7 @@ def forward(self, *inputs: Any) -> Any:
         sample = inputs if len(inputs) > 1 else inputs[0]
         if not (
             has_all(sample, features.BoundingBox)
-            and has_any(sample, PIL.Image.Image, features.Image, is_simple_tensor)
+            and has_any(sample, PIL.Image.Image, features.Image, features.is_simple_tensor)
             and has_any(sample, features.Label, features.OneHotLabel)
         ):
             raise TypeError(

torchvision/prototype/transforms/_meta.py

@@ -7,8 +7,6 @@
 from torchvision.prototype.transforms import functional as F, Transform
 from torchvision.transforms.functional import convert_image_dtype
 
-from ._utils import is_simple_tensor
-
 
 class ConvertBoundingBoxFormat(Transform):
     def __init__(self, format: Union[str, features.BoundingBoxFormat]) -> None:
@@ -34,7 +32,7 @@ def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         if isinstance(inpt, features.Image):
             output = convert_image_dtype(inpt, dtype=self.dtype)
             return features.Image.new_like(inpt, output, dtype=self.dtype)
-        elif is_simple_tensor(inpt):
+        elif features.is_simple_tensor(inpt):
             return convert_image_dtype(inpt, dtype=self.dtype)
         else:
             return inpt

torchvision/prototype/transforms/_type_conversion.py

@@ -8,8 +8,6 @@
 from torchvision.prototype import features
 from torchvision.prototype.transforms import functional as F, Transform
 
-from ._utils import is_simple_tensor
-
 
 class DecodeImage(Transform):
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
@@ -51,7 +49,7 @@ def __init__(self, *, copy: bool = False) -> None:
         self.copy = copy
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        if isinstance(inpt, (features.Image, PIL.Image.Image, np.ndarray)) or is_simple_tensor(inpt):
+        if isinstance(inpt, (features.Image, PIL.Image.Image, np.ndarray)) or features.is_simple_tensor(inpt):
             output = F.to_image_tensor(inpt, copy=self.copy)
             return features.Image(output)
         else:
@@ -68,7 +66,7 @@ def __init__(self, *, mode: Optional[str] = None) -> None:
         self.mode = mode
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        if isinstance(inpt, (features.Image, PIL.Image.Image, np.ndarray)) or is_simple_tensor(inpt):
+        if isinstance(inpt, (features.Image, PIL.Image.Image, np.ndarray)) or features.is_simple_tensor(inpt):
             return F.to_image_pil(inpt, mode=self.mode)
         else:
             return inpt

torchvision/prototype/transforms/_utils.py

@@ -36,7 +36,7 @@ def query_chw(sample: Any) -> Tuple[int, int, int]:
     chws = {
         get_chw(item)
         for item in flat_sample
-        if isinstance(item, (features.Image, PIL.Image.Image)) or is_simple_tensor(item)
+        if isinstance(item, (features.Image, PIL.Image.Image)) or features.is_simple_tensor(item)
     }
     if not chws:
         raise TypeError("No image was found in the sample")
@@ -63,10 +63,3 @@ def has_all(sample: Any, *types_or_checks: Union[Type, Callable[[Any], bool]]) -
         else:
             return False
     return True
-
-
-# TODO: Given that this is not related to pytree / the Transform object, we should probably move it to somewhere else.
-#  One possibility is `functional._utils` so both the functionals and the transforms have proper access to it. We could
-#  also move it `features` since it literally checks for the _Feature type.
-def is_simple_tensor(inpt: Any) -> bool:
-    return isinstance(inpt, torch.Tensor) and not isinstance(inpt, features._Feature)

torchvision/prototype/transforms/functional/_deprecated.py

@@ -6,8 +6,6 @@
 from torchvision.prototype import features
 from torchvision.transforms import functional as _F
 
-from .._utils import is_simple_tensor
-
 
 def to_grayscale(inpt: PIL.Image.Image, num_output_channels: int = 1) -> PIL.Image.Image:
     call = ", num_output_channels=3" if num_output_channels == 3 else ""
@@ -23,7 +21,7 @@ def to_grayscale(inpt: PIL.Image.Image, num_output_channels: int = 1) -> PIL.Ima
 
 
 def rgb_to_grayscale(inpt: Any, num_output_channels: int = 1) -> Any:
-    old_color_space = features.Image.guess_color_space(inpt) if is_simple_tensor(inpt) else None
+    old_color_space = features.Image.guess_color_space(inpt) if features.is_simple_tensor(inpt) else None
 
     call = ", num_output_channels=3" if num_output_channels == 3 else ""
     replacement = (