fix prototype features and functional transforms after review (#5377)

* fix prototype functional transforms after review

* address features review

Author: pmeier

torchvision/prototype/features/_bounding_box.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Tuple, Union, Optional
 
 import torch
@@ -15,7 +17,6 @@ class BoundingBoxFormat(StrEnum):
 
 
 class BoundingBox(Feature):
-    formats = BoundingBoxFormat
     format: BoundingBoxFormat
     image_size: Tuple[int, int]
 
@@ -27,7 +28,7 @@ def __new__(
         device: Optional[torch.device] = None,
         format: Union[BoundingBoxFormat, str],
         image_size: Tuple[int, int],
-    ):
+    ) -> BoundingBox:
         bounding_box = super().__new__(cls, data, dtype=dtype, device=device)
 
         if isinstance(format, str):
@@ -37,7 +38,7 @@ def __new__(
 
         return bounding_box
 
-    def to_format(self, format: Union[str, BoundingBoxFormat]) -> "BoundingBox":
+    def to_format(self, format: Union[str, BoundingBoxFormat]) -> BoundingBox:
         # import at runtime to avoid cyclic imports
         from torchvision.prototype.transforms.functional import convert_bounding_box_format
 

torchvision/prototype/features/_image.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import warnings
 from typing import Any, Optional, Union, Tuple, cast
 
@@ -20,7 +22,6 @@ class ColorSpace(StrEnum):
 
 
 class Image(Feature):
-    color_spaces = ColorSpace
     color_space: ColorSpace
 
     def __new__(
@@ -79,5 +80,5 @@ def guess_color_space(data: torch.Tensor) -> ColorSpace:
     def show(self) -> None:
         to_pil_image(make_grid(self.view(-1, *self.shape[-3:]))).show()
 
-    def draw_bounding_box(self, bounding_box: BoundingBox, **kwargs: Any) -> "Image":
+    def draw_bounding_box(self, bounding_box: BoundingBox, **kwargs: Any) -> Image:
         return Image.new_like(self, draw_bounding_boxes(self, bounding_box.to_format("xyxy").view(-1, 4), **kwargs))

torchvision/prototype/features/_label.py

@@ -26,7 +26,6 @@ def __new__(
 
     @classmethod
     def from_category(cls, category: str, *, categories: Sequence[str]):
-        categories = list(categories)
         return cls(categories.index(category), categories=categories)
 
     def to_categories(self):
@@ -45,7 +44,7 @@ def __new__(
         *,
         dtype: Optional[torch.dtype] = None,
         device: Optional[torch.device] = None,
-        like: Optional["Label"] = None,
+        like: Optional[Label] = None,
         categories: Optional[Sequence[str]] = None,
     ):
         one_hot_label = super().__new__(cls, data, dtype=dtype, device=device)

torchvision/prototype/transforms/__init__.py

@@ -1,8 +1,4 @@
 from . import functional
 from .functional import InterpolationMode  # usort: skip
 
-from ._transform import Transform
-from ._container import Compose, RandomApply, RandomChoice, RandomOrder  # usort: skip
-from ._geometry import Resize, RandomResize, HorizontalFlip, Crop, CenterCrop, RandomCrop
-from ._misc import Identity, Normalize
 from ._presets import CocoEval, ImageNetEval, VocEval, Kinect400Eval, RaftEval