Skip to content

GitLab

Unverified Commit 0de3e5b4 authored by vfdev's avatar vfdev Committed by GitHub
Browse files

[proto] Aligned fill, padding typehints between Features and F (#6616)

parent 841b9a19
Show whitespace changes
Inline Side-by-side
Showing with 45 additions and 52 deletions
torchvision/prototype/features/_bounding_box.py
View file @ 0de3e5b4
......@@ -115,7 +115,7 @@ class BoundingBox(_Feature):
def pad(
self,
padding: Union[int, Sequence[int]],
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
padding_mode: str = "constant",
) -> BoundingBox:
# This cast does Sequence[int] -> List[int] and is required to make mypy happy
......@@ -137,7 +137,7 @@ class BoundingBox(_Feature):
angle: float,
interpolation: InterpolationMode = InterpolationMode.NEAREST,
expand: bool = False,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> BoundingBox:
output = self._F.rotate_bounding_box(
......@@ -165,7 +165,7 @@ class BoundingBox(_Feature):
scale: float,
shear: List[float],
interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> BoundingBox:
output = self._F.affine_bounding_box(
......@@ -184,7 +184,7 @@ class BoundingBox(_Feature):
self,
perspective_coeffs: List[float],
interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> BoundingBox:
output = self._F.perspective_bounding_box(self, self.format, perspective_coeffs)
return BoundingBox.new_like(self, output, dtype=output.dtype)
......@@ -193,7 +193,7 @@ class BoundingBox(_Feature):
self,
displacement: torch.Tensor,
interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = 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
View file @ 0de3e5b4
......@@ -153,8 +153,8 @@ class _Feature(torch.Tensor):
def pad(
self,
padding: Union[int, Sequence[int]],
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
padding: Union[int, List[int]],
fill: Optional[Union[int, float, List[float]]] = None,
padding_mode: str = "constant",
) -> _Feature:
return self
......@@ -164,7 +164,7 @@ class _Feature(torch.Tensor):
angle: float,
interpolation: InterpolationMode = InterpolationMode.NEAREST,
expand: bool = False,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> _Feature:
return self
......@@ -176,7 +176,7 @@ class _Feature(torch.Tensor):
scale: float,
shear: List[float],
interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> _Feature:
return self
......@@ -185,7 +185,7 @@ class _Feature(torch.Tensor):
self,
perspective_coeffs: List[float],
interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> _Feature:
return self
......@@ -193,7 +193,7 @@ class _Feature(torch.Tensor):
self,
displacement: torch.Tensor,
interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> _Feature:
return self
......
torchvision/prototype/features/_image.py
View file @ 0de3e5b4
from __future__ import annotations
import warnings
from typing import Any, cast, List, Optional, Sequence, Tuple, Union
from typing import Any, cast, List, Optional, Tuple, Union
import torch
from torchvision._utils import StrEnum
......@@ -180,16 +180,10 @@ class Image(_Feature):
def pad(
self,
padding: Union[int, Sequence[int]],
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
padding: Union[int, List[int]],
fill: Optional[Union[int, float, List[float]]] = None,
padding_mode: str = "constant",
) -> Image:
# This cast does Sequence[int] -> List[int] and is required to make mypy happy
if not isinstance(padding, int):
padding = list(padding)
fill = self._F._geometry._convert_fill_arg(fill)
output = self._F.pad_image_tensor(self, padding, fill=fill, padding_mode=padding_mode)
return Image.new_like(self, output)
......@@ -198,11 +192,9 @@ class Image(_Feature):
angle: float,
interpolation: InterpolationMode = InterpolationMode.NEAREST,
expand: bool = False,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> Image:
fill = self._F._geometry._convert_fill_arg(fill)
output = self._F._geometry.rotate_image_tensor(
self, angle, interpolation=interpolation, expand=expand, fill=fill, center=center
)
......@@ -215,11 +207,9 @@ class Image(_Feature):
scale: float,
shear: List[float],
interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> Image:
fill = self._F._geometry._convert_fill_arg(fill)
output = self._F._geometry.affine_image_tensor(
self,
angle,
......@@ -236,10 +226,8 @@ class Image(_Feature):
self,
perspective_coeffs: List[float],
interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> Image:
fill = self._F._geometry._convert_fill_arg(fill)
output = self._F._geometry.perspective_image_tensor(
self, perspective_coeffs, interpolation=interpolation, fill=fill
)
......@@ -249,10 +237,8 @@ class Image(_Feature):
self,
displacement: torch.Tensor,
interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> Image:
fill = self._F._geometry._convert_fill_arg(fill)
output = self._F._geometry.elastic_image_tensor(self, displacement, interpolation=interpolation, fill=fill)
return Image.new_like(self, output)
......
torchvision/prototype/features/_mask.py
View file @ 0de3e5b4
from __future__ import annotations
from typing import List, Optional, Sequence, Union
from typing import List, Optional, Union
import torch
from torchvision.transforms import InterpolationMode
......@@ -50,16 +50,10 @@ class Mask(_Feature):
def pad(
self,
padding: Union[int, Sequence[int]],
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
padding: Union[int, List[int]],
fill: Optional[Union[int, float, List[float]]] = None,
padding_mode: str = "constant",
) -> Mask:
# This cast does Sequence[int] -> List[int] and is required to make mypy happy
if not isinstance(padding, int):
padding = list(padding)
fill = self._F._geometry._convert_fill_arg(fill)
output = self._F.pad_mask(self, padding, padding_mode=padding_mode, fill=fill)
return Mask.new_like(self, output)
......@@ -68,10 +62,10 @@ class Mask(_Feature):
angle: float,
interpolation: InterpolationMode = InterpolationMode.NEAREST,
expand: bool = False,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> Mask:
output = self._F.rotate_mask(self, angle, expand=expand, center=center)
output = self._F.rotate_mask(self, angle, expand=expand, center=center, fill=fill)
return Mask.new_like(self, output)
def affine(
......@@ -81,7 +75,7 @@ class Mask(_Feature):
scale: float,
shear: List[float],
interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> Mask:
output = self._F.affine_mask(
......@@ -90,6 +84,7 @@ class Mask(_Feature):
translate=translate,
scale=scale,
shear=shear,
fill=fill,
center=center,
)
return Mask.new_like(self, output)
......@@ -98,16 +93,16 @@ class Mask(_Feature):
self,
perspective_coeffs: List[float],
interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> Mask:
output = self._F.perspective_mask(self, perspective_coeffs)
output = self._F.perspective_mask(self, perspective_coeffs, fill=fill)
return Mask.new_like(self, output)
def elastic(
self,
displacement: torch.Tensor,
interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
fill: Optional[Union[int, float, List[float]]] = None,
) -> Mask:
output = self._F.elastic_mask(self, displacement)
output = self._F.elastic_mask(self, displacement, fill=fill)
return Mask.new_like(self, output, dtype=output.dtype)
torchvision/prototype/transforms/functional/_geometry.py
View file @ 0de3e5b4
......@@ -379,6 +379,7 @@ def affine_mask(
translate: List[float],
scale: float,
shear: List[float],
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> torch.Tensor:
if mask.ndim < 3:
......@@ -394,6 +395,7 @@ def affine_mask(
scale=scale,
shear=shear,
interpolation=InterpolationMode.NEAREST,
fill=fill,
center=center,
)
......@@ -541,6 +543,7 @@ def rotate_mask(
mask: torch.Tensor,
angle: float,
expand: bool = False,
fill: Optional[Union[int, float, List[float]]] = None,
center: Optional[List[float]] = None,
) -> torch.Tensor:
if mask.ndim < 3:
......@@ -554,6 +557,7 @@ def rotate_mask(
angle=angle,
expand=expand,
interpolation=InterpolationMode.NEAREST,
fill=fill,
center=center,
)
......@@ -849,7 +853,11 @@ def perspective_bounding_box(
).view(original_shape)
def perspective_mask(mask: torch.Tensor, perspective_coeffs: List[float]) -> torch.Tensor:
def perspective_mask(
mask: torch.Tensor,
perspective_coeffs: List[float],
fill: Optional[Union[int, float, List[float]]] = None,
) -> torch.Tensor:
if mask.ndim < 3:
mask = mask.unsqueeze(0)
needs_squeeze = True
......@@ -857,7 +865,7 @@ def perspective_mask(mask: torch.Tensor, perspective_coeffs: List[float]) -> tor
needs_squeeze = False
output = perspective_image_tensor(
mask, perspective_coeffs=perspective_coeffs, interpolation=InterpolationMode.NEAREST
mask, perspective_coeffs=perspective_coeffs, interpolation=InterpolationMode.NEAREST, fill=fill
)
if needs_squeeze:
......@@ -944,14 +952,18 @@ def elastic_bounding_box(
).view(original_shape)
def elastic_mask(mask: torch.Tensor, displacement: torch.Tensor) -> torch.Tensor:
def elastic_mask(
mask: torch.Tensor,
displacement: torch.Tensor,
fill: Optional[Union[int, float, List[float]]] = None,
) -> torch.Tensor:
if mask.ndim < 3:
mask = mask.unsqueeze(0)
needs_squeeze = True
else:
needs_squeeze = False
output = elastic_image_tensor(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST)
output = elastic_image_tensor(mask, displacement=displacement, interpolation=InterpolationMode.NEAREST, fill=fill)
if needs_squeeze:
output = output.squeeze(0)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment