int/float support for fill in pad op

Author: vfdev-5

test/test_prototype_transforms_functional.py

@@ -431,7 +431,7 @@ def pad_image_tensor():
     for image, padding, fill, padding_mode in itertools.product(
         make_images(),
         [[1], [1, 1], [1, 1, 2, 2]],  # padding
-        [12],  # fill
+        [12, 12.0],  # fill
         ["constant", "symmetric", "edge", "reflect"],  # padding mode,
     ):
         yield SampleInput(image, padding=padding, fill=fill, padding_mode=padding_mode)

torchvision/prototype/features/_bounding_box.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, List, Tuple, Union, Optional
+from typing import Any, List, Tuple, Union, Optional, Sequence
 
 import torch
 from torchvision._utils import StrEnum
@@ -127,7 +127,9 @@ def resized_crop(
         image_size = (size[0], size[0]) if len(size) == 1 else (size[0], size[1])
         return BoundingBox.new_like(self, output, image_size=image_size, dtype=output.dtype)
 
-    def pad(self, padding: List[int], fill: int = 0, padding_mode: str = "constant") -> BoundingBox:
+    def pad(
+        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+    ) -> BoundingBox:
         from torchvision.prototype.transforms import functional as _F
 
         if padding_mode not in ["constant"]:

torchvision/prototype/features/_feature.py

@@ -119,7 +119,9 @@ def resized_crop(
     ) -> Any:
         return self
 
-    def pad(self, padding: List[int], fill: Union[float, Sequence[float]] = 0, padding_mode: str = "constant") -> Any:
+    def pad(
+        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+    ) -> Any:
         return self
 
     def rotate(

torchvision/prototype/features/_image.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, List, Optional, Union, Tuple, cast
+from typing import Any, List, Optional, Union, Sequence, Tuple, cast
 
 import torch
 from torchvision._utils import StrEnum
@@ -163,7 +163,9 @@ def resized_crop(
         )
         return Image.new_like(self, output)
 
-    def pad(self, padding: List[int], fill: Union[float, List[float]] = 0.0, padding_mode: str = "constant") -> Image:
+    def pad(
+        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+    ) -> Image:
         from torchvision.prototype.transforms import functional as _F
 
         # PyTorch's pad supports only scalars on fill. So we need to overwrite the colour

torchvision/prototype/features/_segmentation_mask.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Tuple, List, Optional
+from typing import Tuple, List, Optional, Union, Sequence
 
 import torch
 from torchvision.transforms import InterpolationMode
@@ -60,7 +60,9 @@ def resized_crop(
         output = _F.resized_crop_segmentation_mask(self, top, left, height, width, size=size)
         return SegmentationMask.new_like(self, output)
 
-    def pad(self, padding: List[int], fill: int = 0, padding_mode: str = "constant") -> SegmentationMask:
+    def pad(
+        self, padding: List[int], fill: Union[int, float, Sequence[float]] = 0, padding_mode: str = "constant"
+    ) -> SegmentationMask:
         from torchvision.prototype.transforms import functional as _F
 
         output = _F.pad_segmentation_mask(self, padding, padding_mode=padding_mode)

torchvision/prototype/transforms/functional/_geometry.py

@@ -507,7 +507,7 @@ def rotate(
 
 
 def pad_image_tensor(
-    img: torch.Tensor, padding: List[int], fill: int = 0, padding_mode: str = "constant"
+    img: torch.Tensor, padding: List[int], fill: Union[int, float] = 0, padding_mode: str = "constant"
 ) -> torch.Tensor:
     num_masks, height, width = img.shape[-3:]
     extra_dims = img.shape[:-3]
@@ -522,8 +522,11 @@ def pad_image_tensor(
 
 # TODO: This should be removed once pytorch pad supports non-scalar padding values
 def _pad_with_vector_fill(
-    img: torch.Tensor, padding: List[int], fill: Union[float, List[float]] = 0.0, padding_mode: str = "constant"
-):
+    img: torch.Tensor,
+    padding: List[int],
+    fill: Sequence[float] = [0.0],
+    padding_mode: str = "constant",
+) -> torch.Tensor:
     if padding_mode != "constant":
         raise ValueError(f"Padding mode '{padding_mode}' is not supported if fill is not scalar")
 
@@ -573,7 +576,7 @@ def pad_bounding_box(
 
 
 def pad(
-    inpt: Any, padding: List[int], fill: Union[float, Sequence[float]] = 0.0, padding_mode: str = "constant"
+    inpt: Any, padding: List[int], fill: Union[int, float, Sequence[float]] = 0.0, padding_mode: str = "constant"
 ) -> Any:
     if isinstance(inpt, features._Feature):
         return inpt.pad(padding, fill=fill, padding_mode=padding_mode)