Nits and TODOs.

Author: datumbox

torchvision/prototype/transforms/_augment.py

@@ -99,6 +99,7 @@ def _transform(self, inpt: features.ImageOrVideoType, params: Dict[str, Any]) ->
         return inpt
 
 
+# TODO: Add support for Video: https://github.com/pytorch/vision/issues/6731
 class _BaseMixupCutmix(_RandomApplyTransform):
     def __init__(self, alpha: float, p: float = 0.5) -> None:
         super().__init__(p=p)

torchvision/prototype/transforms/_auto_augment.py

@@ -521,7 +521,7 @@ def forward(self, *inputs: Any) -> Any:
         mix = mix.view(orig_dims).to(dtype=image_or_video.dtype)
 
         if isinstance(orig_image_or_video, (features.Image, features.Video)):
-            mix = type(orig_image_or_video).wrap_like(orig_image_or_video, mix)  # type: ignore[arg-type]
+            mix = orig_image_or_video.wrap_like(orig_image_or_video, mix)  # type: ignore[arg-type]
         elif isinstance(orig_image_or_video, PIL.Image.Image):
             mix = F.to_image_pil(mix)
 

torchvision/prototype/transforms/_color.py

@@ -119,7 +119,7 @@ def _permute_channels(
         output = inpt[..., permutation, :, :]
 
         if isinstance(inpt, (features.Image, features.Video)):
-            output = type(inpt).wrap_like(inpt, output, color_space=features.ColorSpace.OTHER)  # type: ignore[arg-type]
+            output = inpt.wrap_like(inpt, output, color_space=features.ColorSpace.OTHER)  # type: ignore[arg-type]
 
         elif isinstance(inpt, PIL.Image.Image):
             output = F.to_image_pil(output)

torchvision/prototype/transforms/_deprecated.py

@@ -55,7 +55,7 @@ def __init__(self, num_output_channels: Literal[1, 3] = 1) -> None:
     def _transform(self, inpt: features.ImageOrVideoType, params: Dict[str, Any]) -> features.ImageOrVideoType:
         output = _F.rgb_to_grayscale(inpt, num_output_channels=self.num_output_channels)
         if isinstance(inpt, (features.Image, features.Video)):
-            output = type(inpt).wrap_like(inpt, output, color_space=features.ColorSpace.GRAY)  # type: ignore[arg-type]
+            output = inpt.wrap_like(inpt, output, color_space=features.ColorSpace.GRAY)  # type: ignore[arg-type]
         return output
 
 
@@ -84,5 +84,5 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
     def _transform(self, inpt: features.ImageOrVideoType, params: Dict[str, Any]) -> features.ImageOrVideoType:
         output = _F.rgb_to_grayscale(inpt, num_output_channels=params["num_input_channels"])
         if isinstance(inpt, (features.Image, features.Video)):
-            output = type(inpt).wrap_like(inpt, output, color_space=features.ColorSpace.GRAY)  # type: ignore[arg-type]
+            output = inpt.wrap_like(inpt, output, color_space=features.ColorSpace.GRAY)  # type: ignore[arg-type]
         return output

torchvision/prototype/transforms/_misc.py

@@ -140,6 +140,7 @@ def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         return F.gaussian_blur(inpt, self.kernel_size, **params)
 
 
+# TODO: Enhance as described at https://github.com/pytorch/vision/issues/6697
 class ToDtype(Lambda):
     def __init__(self, dtype: torch.dtype, *types: Type) -> None:
         self.dtype = dtype

torchvision/prototype/transforms/functional/_augment.py

@@ -35,7 +35,7 @@ def erase(
     if isinstance(inpt, torch.Tensor):
         output = erase_image_tensor(inpt, i=i, j=j, h=h, w=w, v=v, inplace=inplace)
         if not torch.jit.is_scripting() and isinstance(inpt, (features.Image, features.Video)):
-            output = type(inpt).wrap_like(inpt, output)  # type: ignore[arg-type]
+            output = inpt.wrap_like(inpt, output)  # type: ignore[arg-type]
         return output
     else:  # isinstance(inpt, PIL.Image.Image):
         return erase_image_pil(inpt, i=i, j=j, h=h, w=w, v=v, inplace=inplace)

torchvision/prototype/transforms/functional/_geometry.py

@@ -1424,8 +1424,7 @@ def five_crop(
     if isinstance(inpt, torch.Tensor):
         output = five_crop_image_tensor(inpt, size)
         if not torch.jit.is_scripting() and isinstance(inpt, (features.Image, features.Video)):
-            cls = type(inpt)
-            output = tuple(cls.wrap_like(inpt, item) for item in output)  # type: ignore[assignment,arg-type]
+            output = (inpt.wrap_like(inpt, item) for item in output)  # type: ignore[arg-type]
         return output
     else:  # isinstance(inpt, PIL.Image.Image):
         return five_crop_image_pil(inpt, size)
@@ -1468,8 +1467,7 @@ def ten_crop(
     if isinstance(inpt, torch.Tensor):
         output = ten_crop_image_tensor(inpt, size, vertical_flip=vertical_flip)
         if not torch.jit.is_scripting() and isinstance(inpt, (features.Image, features.Video)):
-            cls = type(inpt)
-            output = [cls.wrap_like(inpt, item) for item in output]  # type: ignore[arg-type]
+            output = [inpt.wrap_like(inpt, item) for item in output]  # type: ignore[arg-type]
         return output
     else:  # isinstance(inpt, PIL.Image.Image):
         return ten_crop_image_pil(inpt, size, vertical_flip=vertical_flip)