Skip to content

GitLab

Unverified Commit 37618552 authored by Philip Meier's avatar Philip Meier Committed by GitHub
Browse files

Unwrap features before passing them into a kernel (#6807) 

* unwrap features before calling the kernels

* revert double unwrapping

* cleanup

* remove debug raise

* more cleanup
parent d0de55db
Hide whitespace changes
Inline Side-by-side
Showing with 142 additions and 79 deletions
torchvision/prototype/features/_bounding_box.py
View file @ 37618552
...@@ -66,15 +66,23 @@ class BoundingBox(_Feature): ...@@ -66,15 +66,23 @@ class BoundingBox(_Feature):
format = BoundingBoxFormat.from_str(format.upper()) format = BoundingBoxFormat.from_str(format.upper())
return BoundingBox.wrap_like( return BoundingBox.wrap_like(
self, self._F.convert_format_bounding_box(self, old_format=self.format, new_format=format), format=format self,
self._F.convert_format_bounding_box(
self.as_subclass(torch.Tensor), old_format=self.format, new_format=format
),
format=format,
) )
def horizontal_flip(self) -> BoundingBox: def horizontal_flip(self) -> BoundingBox:
output = self._F.horizontal_flip_bounding_box(self, format=self.format, spatial_size=self.spatial_size) output = self._F.horizontal_flip_bounding_box(
self.as_subclass(torch.Tensor), format=self.format, spatial_size=self.spatial_size
)
return BoundingBox.wrap_like(self, output) return BoundingBox.wrap_like(self, output)
def vertical_flip(self) -> BoundingBox: def vertical_flip(self) -> BoundingBox:
output = self._F.vertical_flip_bounding_box(self, format=self.format, spatial_size=self.spatial_size) output = self._F.vertical_flip_bounding_box(
self.as_subclass(torch.Tensor), format=self.format, spatial_size=self.spatial_size
)
return BoundingBox.wrap_like(self, output) return BoundingBox.wrap_like(self, output)
def resize( # type: ignore[override] def resize( # type: ignore[override]
...@@ -85,19 +93,19 @@ class BoundingBox(_Feature): ...@@ -85,19 +93,19 @@ class BoundingBox(_Feature):
antialias: bool = False, antialias: bool = False,
) -> BoundingBox: ) -> BoundingBox:
output, spatial_size = self._F.resize_bounding_box( output, spatial_size = self._F.resize_bounding_box(
self, spatial_size=self.spatial_size, size=size, max_size=max_size self.as_subclass(torch.Tensor), spatial_size=self.spatial_size, size=size, max_size=max_size
) )
return BoundingBox.wrap_like(self, output, spatial_size=spatial_size) return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
def crop(self, top: int, left: int, height: int, width: int) -> BoundingBox: def crop(self, top: int, left: int, height: int, width: int) -> BoundingBox:
output, spatial_size = self._F.crop_bounding_box( output, spatial_size = self._F.crop_bounding_box(
self, self.format, top=top, left=left, height=height, width=width self.as_subclass(torch.Tensor), self.format, top=top, left=left, height=height, width=width
) )
return BoundingBox.wrap_like(self, output, spatial_size=spatial_size) return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
def center_crop(self, output_size: List[int]) -> BoundingBox: def center_crop(self, output_size: List[int]) -> BoundingBox:
output, spatial_size = self._F.center_crop_bounding_box( output, spatial_size = self._F.center_crop_bounding_box(
self, format=self.format, spatial_size=self.spatial_size, output_size=output_size self.as_subclass(torch.Tensor), format=self.format, spatial_size=self.spatial_size, output_size=output_size
) )
return BoundingBox.wrap_like(self, output, spatial_size=spatial_size) return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
...@@ -111,7 +119,9 @@ class BoundingBox(_Feature): ...@@ -111,7 +119,9 @@ class BoundingBox(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
antialias: bool = False, antialias: bool = False,
) -> BoundingBox: ) -> BoundingBox:
output, spatial_size = self._F.resized_crop_bounding_box(self, self.format, top, left, height, width, size=size) output, spatial_size = self._F.resized_crop_bounding_box(
self.as_subclass(torch.Tensor), self.format, top, left, height, width, size=size
)
return BoundingBox.wrap_like(self, output, spatial_size=spatial_size) return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
def pad( def pad(
...@@ -121,7 +131,11 @@ class BoundingBox(_Feature): ...@@ -121,7 +131,11 @@ class BoundingBox(_Feature):
padding_mode: str = "constant", padding_mode: str = "constant",
) -> BoundingBox: ) -> BoundingBox:
output, spatial_size = self._F.pad_bounding_box( output, spatial_size = self._F.pad_bounding_box(
self, format=self.format, spatial_size=self.spatial_size, padding=padding, padding_mode=padding_mode self.as_subclass(torch.Tensor),
format=self.format,
spatial_size=self.spatial_size,
padding=padding,
padding_mode=padding_mode,
) )
return BoundingBox.wrap_like(self, output, spatial_size=spatial_size) return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
...@@ -134,7 +148,12 @@ class BoundingBox(_Feature): ...@@ -134,7 +148,12 @@ class BoundingBox(_Feature):
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> BoundingBox: ) -> BoundingBox:
output, spatial_size = self._F.rotate_bounding_box( output, spatial_size = self._F.rotate_bounding_box(
self, format=self.format, spatial_size=self.spatial_size, angle=angle, expand=expand, center=center self.as_subclass(torch.Tensor),
format=self.format,
spatial_size=self.spatial_size,
angle=angle,
expand=expand,
center=center,
) )
return BoundingBox.wrap_like(self, output, spatial_size=spatial_size) return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
...@@ -149,7 +168,7 @@ class BoundingBox(_Feature): ...@@ -149,7 +168,7 @@ class BoundingBox(_Feature):
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> BoundingBox: ) -> BoundingBox:
output = self._F.affine_bounding_box( output = self._F.affine_bounding_box(
self, self.as_subclass(torch.Tensor),
self.format, self.format,
self.spatial_size, self.spatial_size,
angle, angle,
...@@ -166,7 +185,7 @@ class BoundingBox(_Feature): ...@@ -166,7 +185,7 @@ class BoundingBox(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> BoundingBox: ) -> BoundingBox:
output = self._F.perspective_bounding_box(self, self.format, perspective_coeffs) output = self._F.perspective_bounding_box(self.as_subclass(torch.Tensor), self.format, perspective_coeffs)
return BoundingBox.wrap_like(self, output) return BoundingBox.wrap_like(self, output)
def elastic( def elastic(
...@@ -175,5 +194,5 @@ class BoundingBox(_Feature): ...@@ -175,5 +194,5 @@ class BoundingBox(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> BoundingBox: ) -> BoundingBox:
output = self._F.elastic_bounding_box(self, self.format, displacement) output = self._F.elastic_bounding_box(self.as_subclass(torch.Tensor), self.format, displacement)
return BoundingBox.wrap_like(self, output) return BoundingBox.wrap_like(self, output)
torchvision/prototype/features/_image.py
View file @ 37618552
...@@ -117,17 +117,17 @@ class Image(_Feature): ...@@ -117,17 +117,17 @@ class Image(_Feature):
return Image.wrap_like( return Image.wrap_like(
self, self,
self._F.convert_color_space_image_tensor( self._F.convert_color_space_image_tensor(
self, old_color_space=self.color_space, new_color_space=color_space, copy=copy self.as_subclass(torch.Tensor), old_color_space=self.color_space, new_color_space=color_space, copy=copy
), ),
color_space=color_space, color_space=color_space,
) )
def horizontal_flip(self) -> Image: def horizontal_flip(self) -> Image:
output = self._F.horizontal_flip_image_tensor(self) output = self._F.horizontal_flip_image_tensor(self.as_subclass(torch.Tensor))
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def vertical_flip(self) -> Image: def vertical_flip(self) -> Image:
output = self._F.vertical_flip_image_tensor(self) output = self._F.vertical_flip_image_tensor(self.as_subclass(torch.Tensor))
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def resize( # type: ignore[override] def resize( # type: ignore[override]
...@@ -138,16 +138,16 @@ class Image(_Feature): ...@@ -138,16 +138,16 @@ class Image(_Feature):
antialias: bool = False, antialias: bool = False,
) -> Image: ) -> Image:
output = self._F.resize_image_tensor( output = self._F.resize_image_tensor(
self, size, interpolation=interpolation, max_size=max_size, antialias=antialias self.as_subclass(torch.Tensor), size, interpolation=interpolation, max_size=max_size, antialias=antialias
) )
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def crop(self, top: int, left: int, height: int, width: int) -> Image: def crop(self, top: int, left: int, height: int, width: int) -> Image:
output = self._F.crop_image_tensor(self, top, left, height, width) output = self._F.crop_image_tensor(self.as_subclass(torch.Tensor), top, left, height, width)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def center_crop(self, output_size: List[int]) -> Image: def center_crop(self, output_size: List[int]) -> Image:
output = self._F.center_crop_image_tensor(self, output_size=output_size) output = self._F.center_crop_image_tensor(self.as_subclass(torch.Tensor), output_size=output_size)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def resized_crop( def resized_crop(
...@@ -161,7 +161,14 @@ class Image(_Feature): ...@@ -161,7 +161,14 @@ class Image(_Feature):
antialias: bool = False, antialias: bool = False,
) -> Image: ) -> Image:
output = self._F.resized_crop_image_tensor( output = self._F.resized_crop_image_tensor(
self, top, left, height, width, size=list(size), interpolation=interpolation, antialias=antialias self.as_subclass(torch.Tensor),
top,
left,
height,
width,
size=list(size),
interpolation=interpolation,
antialias=antialias,
) )
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
...@@ -171,7 +178,7 @@ class Image(_Feature): ...@@ -171,7 +178,7 @@ class Image(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
padding_mode: str = "constant", padding_mode: str = "constant",
) -> Image: ) -> Image:
output = self._F.pad_image_tensor(self, padding, fill=fill, padding_mode=padding_mode) output = self._F.pad_image_tensor(self.as_subclass(torch.Tensor), padding, fill=fill, padding_mode=padding_mode)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def rotate( def rotate(
...@@ -182,8 +189,8 @@ class Image(_Feature): ...@@ -182,8 +189,8 @@ class Image(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> Image: ) -> Image:
output = self._F._geometry.rotate_image_tensor( output = self._F.rotate_image_tensor(
self, angle, interpolation=interpolation, expand=expand, fill=fill, center=center self.as_subclass(torch.Tensor), angle, interpolation=interpolation, expand=expand, fill=fill, center=center
) )
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
...@@ -197,8 +204,8 @@ class Image(_Feature): ...@@ -197,8 +204,8 @@ class Image(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> Image: ) -> Image:
output = self._F._geometry.affine_image_tensor( output = self._F.affine_image_tensor(
self, self.as_subclass(torch.Tensor),
angle, angle,
translate=translate, translate=translate,
scale=scale, scale=scale,
...@@ -215,8 +222,8 @@ class Image(_Feature): ...@@ -215,8 +222,8 @@ class Image(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> Image: ) -> Image:
output = self._F._geometry.perspective_image_tensor( output = self._F.perspective_image_tensor(
self, perspective_coeffs, interpolation=interpolation, fill=fill self.as_subclass(torch.Tensor), perspective_coeffs, interpolation=interpolation, fill=fill
) )
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
...@@ -226,55 +233,65 @@ class Image(_Feature): ...@@ -226,55 +233,65 @@ class Image(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> Image: ) -> Image:
output = self._F._geometry.elastic_image_tensor(self, displacement, interpolation=interpolation, fill=fill) output = self._F.elastic_image_tensor(
self.as_subclass(torch.Tensor), displacement, interpolation=interpolation, fill=fill
)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def adjust_brightness(self, brightness_factor: float) -> Image: def adjust_brightness(self, brightness_factor: float) -> Image:
output = self._F.adjust_brightness_image_tensor(self, brightness_factor=brightness_factor) output = self._F.adjust_brightness_image_tensor(
self.as_subclass(torch.Tensor), brightness_factor=brightness_factor
)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def adjust_saturation(self, saturation_factor: float) -> Image: def adjust_saturation(self, saturation_factor: float) -> Image:
output = self._F.adjust_saturation_image_tensor(self, saturation_factor=saturation_factor) output = self._F.adjust_saturation_image_tensor(
self.as_subclass(torch.Tensor), saturation_factor=saturation_factor
)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def adjust_contrast(self, contrast_factor: float) -> Image: def adjust_contrast(self, contrast_factor: float) -> Image:
output = self._F.adjust_contrast_image_tensor(self, contrast_factor=contrast_factor) output = self._F.adjust_contrast_image_tensor(self.as_subclass(torch.Tensor), contrast_factor=contrast_factor)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def adjust_sharpness(self, sharpness_factor: float) -> Image: def adjust_sharpness(self, sharpness_factor: float) -> Image:
output = self._F.adjust_sharpness_image_tensor(self, sharpness_factor=sharpness_factor) output = self._F.adjust_sharpness_image_tensor(
self.as_subclass(torch.Tensor), sharpness_factor=sharpness_factor
)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def adjust_hue(self, hue_factor: float) -> Image: def adjust_hue(self, hue_factor: float) -> Image:
output = self._F.adjust_hue_image_tensor(self, hue_factor=hue_factor) output = self._F.adjust_hue_image_tensor(self.as_subclass(torch.Tensor), hue_factor=hue_factor)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def adjust_gamma(self, gamma: float, gain: float = 1) -> Image: def adjust_gamma(self, gamma: float, gain: float = 1) -> Image:
output = self._F.adjust_gamma_image_tensor(self, gamma=gamma, gain=gain) output = self._F.adjust_gamma_image_tensor(self.as_subclass(torch.Tensor), gamma=gamma, gain=gain)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def posterize(self, bits: int) -> Image: def posterize(self, bits: int) -> Image:
output = self._F.posterize_image_tensor(self, bits=bits) output = self._F.posterize_image_tensor(self.as_subclass(torch.Tensor), bits=bits)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def solarize(self, threshold: float) -> Image: def solarize(self, threshold: float) -> Image:
output = self._F.solarize_image_tensor(self, threshold=threshold) output = self._F.solarize_image_tensor(self.as_subclass(torch.Tensor), threshold=threshold)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def autocontrast(self) -> Image: def autocontrast(self) -> Image:
output = self._F.autocontrast_image_tensor(self) output = self._F.autocontrast_image_tensor(self.as_subclass(torch.Tensor))
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def equalize(self) -> Image: def equalize(self) -> Image:
output = self._F.equalize_image_tensor(self) output = self._F.equalize_image_tensor(self.as_subclass(torch.Tensor))
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def invert(self) -> Image: def invert(self) -> Image:
output = self._F.invert_image_tensor(self) output = self._F.invert_image_tensor(self.as_subclass(torch.Tensor))
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Image: def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Image:
output = self._F.gaussian_blur_image_tensor(self, kernel_size=kernel_size, sigma=sigma) output = self._F.gaussian_blur_image_tensor(
self.as_subclass(torch.Tensor), kernel_size=kernel_size, sigma=sigma
)
return Image.wrap_like(self, output) return Image.wrap_like(self, output)
......
torchvision/prototype/features/_mask.py
View file @ 37618552
...@@ -37,11 +37,11 @@ class Mask(_Feature): ...@@ -37,11 +37,11 @@ class Mask(_Feature):
return cast(Tuple[int, int], tuple(self.shape[-2:])) return cast(Tuple[int, int], tuple(self.shape[-2:]))
def horizontal_flip(self) -> Mask: def horizontal_flip(self) -> Mask:
output = self._F.horizontal_flip_mask(self) output = self._F.horizontal_flip_mask(self.as_subclass(torch.Tensor))
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def vertical_flip(self) -> Mask: def vertical_flip(self) -> Mask:
output = self._F.vertical_flip_mask(self) output = self._F.vertical_flip_mask(self.as_subclass(torch.Tensor))
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def resize( # type: ignore[override] def resize( # type: ignore[override]
...@@ -51,15 +51,15 @@ class Mask(_Feature): ...@@ -51,15 +51,15 @@ class Mask(_Feature):
max_size: Optional[int] = None, max_size: Optional[int] = None,
antialias: bool = False, antialias: bool = False,
) -> Mask: ) -> Mask:
output = self._F.resize_mask(self, size, max_size=max_size) output = self._F.resize_mask(self.as_subclass(torch.Tensor), size, max_size=max_size)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def crop(self, top: int, left: int, height: int, width: int) -> Mask: def crop(self, top: int, left: int, height: int, width: int) -> Mask:
output = self._F.crop_mask(self, top, left, height, width) output = self._F.crop_mask(self.as_subclass(torch.Tensor), top, left, height, width)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def center_crop(self, output_size: List[int]) -> Mask: def center_crop(self, output_size: List[int]) -> Mask:
output = self._F.center_crop_mask(self, output_size=output_size) output = self._F.center_crop_mask(self.as_subclass(torch.Tensor), output_size=output_size)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def resized_crop( def resized_crop(
...@@ -72,7 +72,7 @@ class Mask(_Feature): ...@@ -72,7 +72,7 @@ class Mask(_Feature):
interpolation: InterpolationMode = InterpolationMode.NEAREST, interpolation: InterpolationMode = InterpolationMode.NEAREST,
antialias: bool = False, antialias: bool = False,
) -> Mask: ) -> Mask:
output = self._F.resized_crop_mask(self, top, left, height, width, size=size) output = self._F.resized_crop_mask(self.as_subclass(torch.Tensor), top, left, height, width, size=size)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def pad( def pad(
...@@ -81,7 +81,7 @@ class Mask(_Feature): ...@@ -81,7 +81,7 @@ class Mask(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
padding_mode: str = "constant", padding_mode: str = "constant",
) -> Mask: ) -> Mask:
output = self._F.pad_mask(self, padding, padding_mode=padding_mode, fill=fill) output = self._F.pad_mask(self.as_subclass(torch.Tensor), padding, padding_mode=padding_mode, fill=fill)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def rotate( def rotate(
...@@ -92,7 +92,7 @@ class Mask(_Feature): ...@@ -92,7 +92,7 @@ class Mask(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> Mask: ) -> Mask:
output = self._F.rotate_mask(self, angle, expand=expand, center=center, fill=fill) output = self._F.rotate_mask(self.as_subclass(torch.Tensor), angle, expand=expand, center=center, fill=fill)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def affine( def affine(
...@@ -106,7 +106,7 @@ class Mask(_Feature): ...@@ -106,7 +106,7 @@ class Mask(_Feature):
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> Mask: ) -> Mask:
output = self._F.affine_mask( output = self._F.affine_mask(
self, self.as_subclass(torch.Tensor),
angle, angle,
translate=translate, translate=translate,
scale=scale, scale=scale,
...@@ -122,7 +122,7 @@ class Mask(_Feature): ...@@ -122,7 +122,7 @@ class Mask(_Feature):
interpolation: InterpolationMode = InterpolationMode.NEAREST, interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> Mask: ) -> Mask:
output = self._F.perspective_mask(self, perspective_coeffs, fill=fill) output = self._F.perspective_mask(self.as_subclass(torch.Tensor), perspective_coeffs, fill=fill)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
def elastic( def elastic(
...@@ -131,5 +131,5 @@ class Mask(_Feature): ...@@ -131,5 +131,5 @@ class Mask(_Feature):
interpolation: InterpolationMode = InterpolationMode.NEAREST, interpolation: InterpolationMode = InterpolationMode.NEAREST,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> Mask: ) -> Mask:
output = self._F.elastic_mask(self, displacement, fill=fill) output = self._F.elastic_mask(self.as_subclass(torch.Tensor), displacement, fill=fill)
return Mask.wrap_like(self, output) return Mask.wrap_like(self, output)
torchvision/prototype/features/_video.py
View file @ 37618552
...@@ -73,17 +73,17 @@ class Video(_Feature): ...@@ -73,17 +73,17 @@ class Video(_Feature):
return Video.wrap_like( return Video.wrap_like(
self, self,
self._F.convert_color_space_video( self._F.convert_color_space_video(
self, old_color_space=self.color_space, new_color_space=color_space, copy=copy self.as_subclass(torch.Tensor), old_color_space=self.color_space, new_color_space=color_space, copy=copy
), ),
color_space=color_space, color_space=color_space,
) )
def horizontal_flip(self) -> Video: def horizontal_flip(self) -> Video:
output = self._F.horizontal_flip_video(self) output = self._F.horizontal_flip_video(self.as_subclass(torch.Tensor))
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def vertical_flip(self) -> Video: def vertical_flip(self) -> Video:
output = self._F.vertical_flip_video(self) output = self._F.vertical_flip_video(self.as_subclass(torch.Tensor))
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def resize( # type: ignore[override] def resize( # type: ignore[override]
...@@ -93,15 +93,21 @@ class Video(_Feature): ...@@ -93,15 +93,21 @@ class Video(_Feature):
max_size: Optional[int] = None, max_size: Optional[int] = None,
antialias: bool = False, antialias: bool = False,
) -> Video: ) -> Video:
output = self._F.resize_video(self, size, interpolation=interpolation, max_size=max_size, antialias=antialias) output = self._F.resize_video(
self.as_subclass(torch.Tensor),
size,
interpolation=interpolation,
max_size=max_size,
antialias=antialias,
)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def crop(self, top: int, left: int, height: int, width: int) -> Video: def crop(self, top: int, left: int, height: int, width: int) -> Video:
output = self._F.crop_video(self, top, left, height, width) output = self._F.crop_video(self.as_subclass(torch.Tensor), top, left, height, width)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def center_crop(self, output_size: List[int]) -> Video: def center_crop(self, output_size: List[int]) -> Video:
output = self._F.center_crop_video(self, output_size=output_size) output = self._F.center_crop_video(self.as_subclass(torch.Tensor), output_size=output_size)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def resized_crop( def resized_crop(
...@@ -115,7 +121,14 @@ class Video(_Feature): ...@@ -115,7 +121,14 @@ class Video(_Feature):
antialias: bool = False, antialias: bool = False,
) -> Video: ) -> Video:
output = self._F.resized_crop_video( output = self._F.resized_crop_video(
self, top, left, height, width, size=list(size), interpolation=interpolation, antialias=antialias self.as_subclass(torch.Tensor),
top,
left,
height,
width,
size=list(size),
interpolation=interpolation,
antialias=antialias,
) )
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
...@@ -125,7 +138,7 @@ class Video(_Feature): ...@@ -125,7 +138,7 @@ class Video(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
padding_mode: str = "constant", padding_mode: str = "constant",
) -> Video: ) -> Video:
output = self._F.pad_video(self, padding, fill=fill, padding_mode=padding_mode) output = self._F.pad_video(self.as_subclass(torch.Tensor), padding, fill=fill, padding_mode=padding_mode)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def rotate( def rotate(
...@@ -136,8 +149,8 @@ class Video(_Feature): ...@@ -136,8 +149,8 @@ class Video(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> Video: ) -> Video:
output = self._F._geometry.rotate_video( output = self._F.rotate_video(
self, angle, interpolation=interpolation, expand=expand, fill=fill, center=center self.as_subclass(torch.Tensor), angle, interpolation=interpolation, expand=expand, fill=fill, center=center
) )
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
...@@ -151,8 +164,8 @@ class Video(_Feature): ...@@ -151,8 +164,8 @@ class Video(_Feature):
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
center: Optional[List[float]] = None, center: Optional[List[float]] = None,
) -> Video: ) -> Video:
output = self._F._geometry.affine_video( output = self._F.affine_video(
self, self.as_subclass(torch.Tensor),
angle, angle,
translate=translate, translate=translate,
scale=scale, scale=scale,
...@@ -169,7 +182,9 @@ class Video(_Feature): ...@@ -169,7 +182,9 @@ class Video(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> Video: ) -> Video:
output = self._F._geometry.perspective_video(self, perspective_coeffs, interpolation=interpolation, fill=fill) output = self._F.perspective_video(
self.as_subclass(torch.Tensor), perspective_coeffs, interpolation=interpolation, fill=fill
)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def elastic( def elastic(
...@@ -178,55 +193,57 @@ class Video(_Feature): ...@@ -178,55 +193,57 @@ class Video(_Feature):
interpolation: InterpolationMode = InterpolationMode.BILINEAR, interpolation: InterpolationMode = InterpolationMode.BILINEAR,
fill: FillTypeJIT = None, fill: FillTypeJIT = None,
) -> Video: ) -> Video:
output = self._F._geometry.elastic_video(self, displacement, interpolation=interpolation, fill=fill) output = self._F.elastic_video(
self.as_subclass(torch.Tensor), displacement, interpolation=interpolation, fill=fill
)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def adjust_brightness(self, brightness_factor: float) -> Video: def adjust_brightness(self, brightness_factor: float) -> Video:
output = self._F.adjust_brightness_video(self, brightness_factor=brightness_factor) output = self._F.adjust_brightness_video(self.as_subclass(torch.Tensor), brightness_factor=brightness_factor)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def adjust_saturation(self, saturation_factor: float) -> Video: def adjust_saturation(self, saturation_factor: float) -> Video:
output = self._F.adjust_saturation_video(self, saturation_factor=saturation_factor) output = self._F.adjust_saturation_video(self.as_subclass(torch.Tensor), saturation_factor=saturation_factor)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def adjust_contrast(self, contrast_factor: float) -> Video: def adjust_contrast(self, contrast_factor: float) -> Video:
output = self._F.adjust_contrast_video(self, contrast_factor=contrast_factor) output = self._F.adjust_contrast_video(self.as_subclass(torch.Tensor), contrast_factor=contrast_factor)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def adjust_sharpness(self, sharpness_factor: float) -> Video: def adjust_sharpness(self, sharpness_factor: float) -> Video:
output = self._F.adjust_sharpness_video(self, sharpness_factor=sharpness_factor) output = self._F.adjust_sharpness_video(self.as_subclass(torch.Tensor), sharpness_factor=sharpness_factor)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def adjust_hue(self, hue_factor: float) -> Video: def adjust_hue(self, hue_factor: float) -> Video:
output = self._F.adjust_hue_video(self, hue_factor=hue_factor) output = self._F.adjust_hue_video(self.as_subclass(torch.Tensor), hue_factor=hue_factor)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def adjust_gamma(self, gamma: float, gain: float = 1) -> Video: def adjust_gamma(self, gamma: float, gain: float = 1) -> Video:
output = self._F.adjust_gamma_video(self, gamma=gamma, gain=gain) output = self._F.adjust_gamma_video(self.as_subclass(torch.Tensor), gamma=gamma, gain=gain)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def posterize(self, bits: int) -> Video: def posterize(self, bits: int) -> Video:
output = self._F.posterize_video(self, bits=bits) output = self._F.posterize_video(self.as_subclass(torch.Tensor), bits=bits)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def solarize(self, threshold: float) -> Video: def solarize(self, threshold: float) -> Video:
output = self._F.solarize_video(self, threshold=threshold) output = self._F.solarize_video(self.as_subclass(torch.Tensor), threshold=threshold)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def autocontrast(self) -> Video: def autocontrast(self) -> Video:
output = self._F.autocontrast_video(self) output = self._F.autocontrast_video(self.as_subclass(torch.Tensor))
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def equalize(self) -> Video: def equalize(self) -> Video:
output = self._F.equalize_video(self) output = self._F.equalize_video(self.as_subclass(torch.Tensor))
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def invert(self) -> Video: def invert(self) -> Video:
output = self._F.invert_video(self) output = self._F.invert_video(self.as_subclass(torch.Tensor))
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Video: def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Video:
output = self._F.gaussian_blur_video(self, kernel_size=kernel_size, sigma=sigma) output = self._F.gaussian_blur_video(self.as_subclass(torch.Tensor), kernel_size=kernel_size, sigma=sigma)
return Video.wrap_like(self, output) return Video.wrap_like(self, output)
......
torchvision/prototype/transforms/_meta.py
View file @ 37618552
...@@ -17,7 +17,11 @@ class ConvertBoundingBoxFormat(Transform): ...@@ -17,7 +17,11 @@ class ConvertBoundingBoxFormat(Transform):
self.format = format self.format = format
def _transform(self, inpt: features.BoundingBox, params: Dict[str, Any]) -> features.BoundingBox: def _transform(self, inpt: features.BoundingBox, params: Dict[str, Any]) -> features.BoundingBox:
output = F.convert_format_bounding_box(inpt, old_format=inpt.format, new_format=params["format"]) # We need to unwrap here to avoid unnecessary `__torch_function__` calls,
# since `convert_format_bounding_box` does not have a dispatcher function that would do that for us
output = F.convert_format_bounding_box(
inpt.as_subclass(torch.Tensor), old_format=inpt.format, new_format=params["format"]
)
return features.BoundingBox.wrap_like(inpt, output, format=params["format"]) return features.BoundingBox.wrap_like(inpt, output, format=params["format"])
...@@ -31,7 +35,9 @@ class ConvertImageDtype(Transform): ...@@ -31,7 +35,9 @@ class ConvertImageDtype(Transform):
def _transform( def _transform(
self, inpt: Union[features.TensorImageType, features.TensorVideoType], params: Dict[str, Any] self, inpt: Union[features.TensorImageType, features.TensorVideoType], params: Dict[str, Any]
) -> Union[features.TensorImageType, features.TensorVideoType]: ) -> Union[features.TensorImageType, features.TensorVideoType]:
output = F.convert_image_dtype(inpt, dtype=self.dtype) # TODO: the `inpt.as_subclass(torch.Tensor)` call can be removed as soon as we have a proper dispatcher that
# handles this. See https://github.com/pytorch/vision/pull/6783 for details.
output = F.convert_image_dtype(inpt.as_subclass(torch.Tensor), dtype=self.dtype)
return ( return (
output if features.is_simple_tensor(inpt) else type(inpt).wrap_like(inpt, output) # type: ignore[attr-defined] output if features.is_simple_tensor(inpt) else type(inpt).wrap_like(inpt, output) # type: ignore[attr-defined]
) )
...@@ -70,5 +76,9 @@ class ClampBoundingBoxes(Transform): ...@@ -70,5 +76,9 @@ class ClampBoundingBoxes(Transform):
_transformed_types = (features.BoundingBox,) _transformed_types = (features.BoundingBox,)
def _transform(self, inpt: features.BoundingBox, params: Dict[str, Any]) -> features.BoundingBox: def _transform(self, inpt: features.BoundingBox, params: Dict[str, Any]) -> features.BoundingBox:
output = F.clamp_bounding_box(inpt, format=inpt.format, spatial_size=inpt.spatial_size) # We need to unwrap here to avoid unnecessary `__torch_function__` calls,
# since `clamp_bounding_box` does not have a dispatcher function that would do that for us
output = F.clamp_bounding_box(
inpt.as_subclass(torch.Tensor), format=inpt.format, spatial_size=inpt.spatial_size
)
return features.BoundingBox.wrap_like(inpt, output) return features.BoundingBox.wrap_like(inpt, output)
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