[proto] Added some transformations and fixed type hints

test/test_functional_tensor.py

@@ -955,18 +955,7 @@ def test_adjust_gamma(device, dtype, config, channels):
 
 @pytest.mark.parametrize("device", cpu_and_gpu())
 @pytest.mark.parametrize("dt", [None, torch.float32, torch.float64, torch.float16])
-@pytest.mark.parametrize(
-    "pad",
-    [
-        2,
-        [
-            3,
-        ],
-        [0, 3],
-        (3, 3),
-        [4, 2, 4, 3],
-    ],
-)
+@pytest.mark.parametrize("pad", [2, [3], [0, 3], (3, 3), [4, 2, 4, 3]])
 @pytest.mark.parametrize(
     "config",
     [

test/test_prototype_transforms_functional.py

@@ -317,7 +317,7 @@ def rotate_image_tensor():
         [-87, 15, 90],  # angle
         [True, False],  # expand
         [None, [12, 23]],  # center
-        [None, [128]],  # fill
+        [None, [128], [12.0]],  # fill
     ):
         if center is not None and expand:
             # Skip warning: The provided center argument is ignored if expand is True

test/test_transforms_tensor.py

@@ -291,11 +291,7 @@ def test_center_crop(device, tmpdir):
     scripted_fn(tensor)
 
     # Test torchscript of transforms.CenterCrop with size as [int, ]
-    f = T.CenterCrop(
-        size=[
-            5,
-        ]
-    )
+    f = T.CenterCrop(size=[5])
     scripted_fn = torch.jit.script(f)
     scripted_fn(tensor)
 
@@ -317,17 +313,7 @@ def test_center_crop(device, tmpdir):
         (F.ten_crop, T.TenCrop, 10),
     ],
 )
-@pytest.mark.parametrize(
-    "size",
-    [
-        (5,),
-        [
-            5,
-        ],
-        (4, 5),
-        [4, 5],
-    ],
-)
+@pytest.mark.parametrize("size", [(5,), [5], (4, 5), [4, 5]])
 def test_x_crop(fn, method, out_length, size, device):
     meth_kwargs = fn_kwargs = {"size": size}
     scripted_fn = torch.jit.script(fn)
@@ -509,19 +495,7 @@ def test_random_affine_degrees(device, interpolation, degrees):
 
 @pytest.mark.parametrize("device", cpu_and_gpu())
 @pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
-@pytest.mark.parametrize(
-    "fill",
-    [
-        85,
-        (10, -10, 10),
-        0.7,
-        [0.0, 0.0, 0.0],
-        [
-            1,
-        ],
-        1,
-    ],
-)
+@pytest.mark.parametrize("fill", [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1], 1])
 def test_random_affine_fill(device, interpolation, fill):
     _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, fill=fill)
 
@@ -531,19 +505,7 @@ def test_random_affine_fill(device, interpolation, fill):
 @pytest.mark.parametrize("expand", [True, False])
 @pytest.mark.parametrize("degrees", [45, 35.0, (-45, 45), [-90.0, 90.0]])
 @pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
-@pytest.mark.parametrize(
-    "fill",
-    [
-        85,
-        (10, -10, 10),
-        0.7,
-        [0.0, 0.0, 0.0],
-        [
-            1,
-        ],
-        1,
-    ],
-)
+@pytest.mark.parametrize("fill", [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1], 1])
 def test_random_rotate(device, center, expand, degrees, interpolation, fill):
     tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
     batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
@@ -564,19 +526,7 @@ def test_random_rotate_save(tmpdir):
 @pytest.mark.parametrize("device", cpu_and_gpu())
 @pytest.mark.parametrize("distortion_scale", np.linspace(0.1, 1.0, num=20))
 @pytest.mark.parametrize("interpolation", [NEAREST, BILINEAR])
-@pytest.mark.parametrize(
-    "fill",
-    [
-        85,
-        (10, -10, 10),
-        0.7,
-        [0.0, 0.0, 0.0],
-        [
-            1,
-        ],
-        1,
-    ],
-)
+@pytest.mark.parametrize("fill", [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1], 1])
 def test_random_perspective(device, distortion_scale, interpolation, fill):
     tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
     batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)

torchvision/prototype/features/_bounding_box.py

@@ -128,13 +128,20 @@ def resized_crop(
         return BoundingBox.new_like(self, output, image_size=image_size, dtype=output.dtype)
 
     def pad(
-        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+        self,
+        padding: Union[int, Sequence[int]],
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
+        padding_mode: str = "constant",
     ) -> BoundingBox:
         from torchvision.prototype.transforms import functional as _F
 
         if padding_mode not in ["constant"]:
             raise ValueError(f"Padding mode '{padding_mode}' is not supported with bounding boxes")
 
+        # This cast does Sequence[int] -> List[int] and is required to make mypy happy
+        if not isinstance(padding, int):
+            padding = list(padding)
+
         output = _F.pad_bounding_box(self, padding, format=self.format)
 
         # Update output image size:
@@ -153,7 +160,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> BoundingBox:
         from torchvision.prototype.transforms import functional as _F
@@ -173,7 +180,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> BoundingBox:
         from torchvision.prototype.transforms import functional as _F
@@ -200,12 +207,3 @@ def perspective(
 
         output = _F.perspective_bounding_box(self, self.format, perspective_coeffs)
         return BoundingBox.new_like(self, output, dtype=output.dtype)
-
-    def erase(self, i: int, j: int, h: int, w: int, v: torch.Tensor) -> BoundingBox:
-        raise TypeError("Erase transformation does not support bounding boxes")
-
-    def mixup(self, lam: float) -> BoundingBox:
-        raise TypeError("Mixup transformation does not support bounding boxes")
-
-    def cutmix(self, box: Tuple[int, int, int, int], lam_adjusted: float) -> BoundingBox:
-        raise TypeError("Cutmix transformation does not support bounding boxes")

torchvision/prototype/features/_feature.py

@@ -120,7 +120,10 @@ def resized_crop(
         return self
 
     def pad(
-        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+        self,
+        padding: Union[int, Sequence[int]],
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
+        padding_mode: str = "constant",
     ) -> Any:
         return self
 
@@ -129,7 +132,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> Any:
         return self
@@ -141,7 +144,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> Any:
         return self
@@ -186,12 +189,3 @@ def equalize(self) -> Any:
 
     def invert(self) -> Any:
         return self
-
-    def erase(self, i: int, j: int, h: int, w: int, v: torch.Tensor) -> Any:
-        return self
-
-    def mixup(self, lam: float) -> Any:
-        return self
-
-    def cutmix(self, box: Tuple[int, int, int, int], lam_adjusted: float) -> Any:
-        return self

torchvision/prototype/features/_image.py

@@ -164,10 +164,17 @@ def resized_crop(
         return Image.new_like(self, output)
 
     def pad(
-        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+        self,
+        padding: Union[int, Sequence[int]],
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
+        padding_mode: str = "constant",
     ) -> Image:
         from torchvision.prototype.transforms import functional as _F
 
+        # This cast does Sequence[int] -> List[int] and is required to make mypy happy
+        if not isinstance(padding, int):
+            padding = list(padding)
+
         # PyTorch's pad supports only scalars on fill. So we need to overwrite the colour
         if isinstance(fill, (int, float)):
             output = _F.pad_image_tensor(self, padding, fill=fill, padding_mode=padding_mode)
@@ -183,11 +190,19 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> Image:
         from torchvision.prototype.transforms import functional as _F
 
+        # This cast does Sequence -> List[float] to please mypy and torch.jit.script
+        if not isinstance(fill, (int, float)):
+            fill = [float(v) for v in list(fill)]
+
+        if isinstance(fill, (int, float)):
+            # It is OK to cast int to float as later we use inpt.dtype
+            fill = [float(fill)]
+
         output = _F.rotate_image_tensor(
             self, angle, interpolation=interpolation, expand=expand, fill=fill, center=center
         )
@@ -200,11 +215,19 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> Image:
         from torchvision.prototype.transforms import functional as _F
 
+        # This cast does Sequence -> List[float] to please mypy and torch.jit.script
+        if not isinstance(fill, (int, float)):
+            fill = [float(v) for v in list(fill)]
+
+        if isinstance(fill, (int, float)):
+            # It is OK to cast int to float as later we use inpt.dtype
+            fill = [float(fill)]
+
         output = _F.affine_image_tensor(
             self,
             angle,
@@ -293,25 +316,3 @@ def invert(self) -> Image:
 
         output = _F.invert_image_tensor(self)
         return Image.new_like(self, output)
-
-    def erase(self, i: int, j: int, h: int, w: int, v: torch.Tensor) -> Image:
-        from torchvision.prototype.transforms import functional as _F
-
-        output = _F.erase_image_tensor(self, i, j, h, w, v)
-        return Image.new_like(self, output)
-
-    def mixup(self, lam: float) -> Image:
-        if self.ndim < 4:
-            raise ValueError("Need a batch of images")
-        output = self.clone()
-        output = output.roll(1, -4).mul_(1 - lam).add_(output.mul_(lam))
-        return Image.new_like(self, output)
-
-    def cutmix(self, box: Tuple[int, int, int, int], lam_adjusted: float) -> Image:
-        if self.ndim < 4:
-            raise ValueError("Need a batch of images")
-        x1, y1, x2, y2 = box
-        image_rolled = self.roll(1, -4)
-        output = self.clone()
-        output[..., y1:y2, x1:x2] = image_rolled[..., y1:y2, x1:x2]
-        return Image.new_like(self, output)

torchvision/prototype/features/_label.py

@@ -77,14 +77,3 @@ def new_like(
         return super().new_like(
             other, data, categories=categories if categories is not None else other.categories, **kwargs
         )
-
-    def mixup(self, lam: float) -> OneHotLabel:
-        if self.ndim < 2:
-            raise ValueError("Need a batch of one hot labels")
-        output = self.clone()
-        output = output.roll(1, -2).mul_(1 - lam).add_(output.mul_(lam))
-        return OneHotLabel.new_like(self, output)
-
-    def cutmix(self, box: Tuple[int, int, int, int], lam_adjusted: float) -> OneHotLabel:
-        box  # unused
-        return self.mixup(lam_adjusted)

torchvision/prototype/features/_segmentation_mask.py

@@ -61,10 +61,17 @@ def resized_crop(
         return SegmentationMask.new_like(self, output)
 
     def pad(
-        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+        self,
+        padding: Union[int, Sequence[int]],
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
+        padding_mode: str = "constant",
     ) -> SegmentationMask:
         from torchvision.prototype.transforms import functional as _F
 
+        # This cast does Sequence[int] -> List[int] and is required to make mypy happy
+        if not isinstance(padding, int):
+            padding = list(padding)
+
         output = _F.pad_segmentation_mask(self, padding, padding_mode=padding_mode)
         return SegmentationMask.new_like(self, output)
 
@@ -73,7 +80,7 @@ def rotate(
         angle: float,
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
         expand: bool = False,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> SegmentationMask:
         from torchvision.prototype.transforms import functional as _F
@@ -88,7 +95,7 @@ def affine(
         scale: float,
         shear: List[float],
         interpolation: InterpolationMode = InterpolationMode.NEAREST,
-        fill: Optional[List[float]] = None,
+        fill: Union[int, float, Sequence[int], Sequence[float]] = 0,
         center: Optional[List[float]] = None,
     ) -> SegmentationMask:
         from torchvision.prototype.transforms import functional as _F
@@ -113,12 +120,3 @@ def perspective(
 
         output = _F.perspective_segmentation_mask(self, perspective_coeffs)
         return SegmentationMask.new_like(self, output)
-
-    def erase(self, i: int, j: int, h: int, w: int, v: torch.Tensor) -> SegmentationMask:
-        raise TypeError("Erase transformation does not support segmentation masks")
-
-    def mixup(self, lam: float) -> SegmentationMask:
-        raise TypeError("Mixup transformation does not support segmentation masks")
-
-    def cutmix(self, box: Tuple[int, int, int, int], lam_adjusted: float) -> SegmentationMask:
-        raise TypeError("Cutmix transformation does not support segmentation masks")