remove or privatize functionality in features / datasets / transforms

pmeier · pmeier · commit 149076701839 · 2022-02-11T11:05:10.000+01:00
diff --git a/test/builtin_dataset_mocks.py b/test/builtin_dataset_mocks.py
@@ -19,7 +19,7 @@
 from datasets_utils import make_zip, make_tar, create_image_folder, create_image_file
 from torch.nn.functional import one_hot
 from torch.testing import make_tensor as _make_tensor
-from torchvision.prototype.datasets._api import find
+from torchvision.prototype.datasets._api import _find
 from torchvision.prototype.utils._internal import sequence_to_str
 
 make_tensor = functools.partial(_make_tensor, device="cpu")
@@ -31,7 +31,7 @@
 
 class DatasetMock:
     def __init__(self, name, mock_data_fn):
-        self.dataset = find(name)
+        self.dataset = _find(name)
         self.info = self.dataset.info
         self.name = self.info.name
 
diff --git a/torchvision/prototype/datasets/__init__.py b/torchvision/prototype/datasets/__init__.py
@@ -11,5 +11,5 @@
 from ._home import home
 
 # Load this last, since some parts depend on the above being loaded first
-from ._api import register, list_datasets, info, load  # usort: skip
+from ._api import list_datasets, load  # usort: skip
 from ._folder import from_data_folder, from_image_folder
diff --git a/torchvision/prototype/datasets/_api.py b/torchvision/prototype/datasets/_api.py
@@ -24,7 +24,7 @@ def list_datasets() -> List[str]:
     return sorted(DATASETS.keys())
 
 
-def find(name: str) -> Dataset:
+def _find(name: str) -> Dataset:
     name = name.lower()
     try:
         return DATASETS[name]
@@ -42,7 +42,7 @@ def find(name: str) -> Dataset:
 
 
 def info(name: str) -> DatasetInfo:
-    return find(name).info
+    return _find(name).info
 
 
 def load(
@@ -51,7 +51,7 @@ def load(
     skip_integrity_check: bool = False,
     **options: Any,
 ) -> IterDataPipe[Dict[str, Any]]:
-    dataset = find(name)
+    dataset = _find(name)
 
     config = dataset.info.make_config(**options)
     root = os.path.join(home(), dataset.name)
diff --git a/torchvision/prototype/datasets/_builtin/caltech.py b/torchvision/prototype/datasets/_builtin/caltech.py
@@ -17,7 +17,7 @@
     OnlineResource,
 )
 from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE, read_mat, hint_sharding, hint_shuffling
-from torchvision.prototype.features import Label, BoundingBox, Feature, EncodedImage
+from torchvision.prototype.features import Label, BoundingBox, _Feature, EncodedImage
 
 
 class Caltech101(Dataset):
@@ -95,7 +95,7 @@ def _prepare_sample(
             bounding_box=BoundingBox(
                 ann["box_coord"].astype(np.int64).squeeze()[[2, 0, 3, 1]], format="xyxy", image_size=image.image_size
             ),
-            contour=Feature(ann["obj_contour"].T),
+            contour=_Feature(ann["obj_contour"].T),
         )
 
     def _make_datapipe(
diff --git a/torchvision/prototype/datasets/_builtin/celeba.py b/torchvision/prototype/datasets/_builtin/celeba.py
@@ -23,7 +23,7 @@
     hint_sharding,
     hint_shuffling,
 )
-from torchvision.prototype.features import EncodedImage, Feature, Label, BoundingBox
+from torchvision.prototype.features import EncodedImage, _Feature, Label, BoundingBox
 
 
 csv.register_dialect("celeba", delimiter=" ", skipinitialspace=True)
@@ -140,7 +140,7 @@ def _prepare_sample(
                 image_size=image.image_size,
             ),
             landmarks={
-                landmark: Feature((int(landmarks[f"{landmark}_x"]), int(landmarks[f"{landmark}_y"])))
+                landmark: _Feature((int(landmarks[f"{landmark}_x"]), int(landmarks[f"{landmark}_y"])))
                 for landmark in {key[:-2] for key in landmarks.keys()}
             },
         )
diff --git a/torchvision/prototype/datasets/_builtin/coco.py b/torchvision/prototype/datasets/_builtin/coco.py
@@ -31,7 +31,7 @@
     hint_sharding,
     hint_shuffling,
 )
-from torchvision.prototype.features import BoundingBox, Label, Feature, EncodedImage
+from torchvision.prototype.features import BoundingBox, Label, _Feature, EncodedImage
 from torchvision.prototype.utils._internal import FrozenMapping
 
 
@@ -95,16 +95,16 @@ def _decode_instances_anns(self, anns: List[Dict[str, Any]], image_meta: Dict[st
         labels = [ann["category_id"] for ann in anns]
         return dict(
             # TODO: create a segmentation feature
-            segmentations=Feature(
+            segmentations=_Feature(
                 torch.stack(
                     [
                         self._segmentation_to_mask(ann["segmentation"], is_crowd=ann["iscrowd"], image_size=image_size)
                         for ann in anns
                     ]
                 )
             ),
-            areas=Feature([ann["area"] for ann in anns]),
-            crowds=Feature([ann["iscrowd"] for ann in anns], dtype=torch.bool),
+            areas=_Feature([ann["area"] for ann in anns]),
+            crowds=_Feature([ann["iscrowd"] for ann in anns], dtype=torch.bool),
             bounding_boxes=BoundingBox(
                 [ann["bbox"] for ann in anns],
                 format="xywh",
diff --git a/torchvision/prototype/datasets/_builtin/cub200.py b/torchvision/prototype/datasets/_builtin/cub200.py
@@ -29,7 +29,7 @@
     path_comparator,
     path_accessor,
 )
-from torchvision.prototype.features import Label, BoundingBox, Feature, EncodedImage
+from torchvision.prototype.features import Label, BoundingBox, _Feature, EncodedImage
 
 csv.register_dialect("cub200", delimiter=" ")
 
@@ -131,7 +131,7 @@ def _2010_prepare_ann(self, data: Tuple[str, Tuple[str, BinaryIO]], image_size:
                 format="xyxy",
                 image_size=image_size,
             ),
-            segmentation=Feature(content["seg"]),
+            segmentation=_Feature(content["seg"]),
         )
 
     def _prepare_sample(
diff --git a/torchvision/prototype/datasets/_builtin/sbd.py b/torchvision/prototype/datasets/_builtin/sbd.py
@@ -27,7 +27,7 @@
     hint_sharding,
     hint_shuffling,
 )
-from torchvision.prototype.features import Feature, EncodedImage
+from torchvision.prototype.features import _Feature, EncodedImage
 
 
 class SBD(Dataset):
@@ -81,8 +81,8 @@ def _prepare_sample(self, data: Tuple[Tuple[Any, Tuple[str, BinaryIO]], Tuple[st
             image=EncodedImage.from_file(image_buffer),
             ann_path=ann_path,
             # the boundaries are stored in sparse CSC format, which is not supported by PyTorch
-            boundaries=Feature(np.stack([raw_boundary.toarray() for raw_boundary in anns["Boundaries"].item()])),
-            segmentation=Feature(anns["Segmentation"].item()),
+            boundaries=_Feature(np.stack([raw_boundary.toarray() for raw_boundary in anns["Boundaries"].item()])),
+            segmentation=_Feature(anns["Segmentation"].item()),
         )
 
     def _make_datapipe(
diff --git a/torchvision/prototype/datasets/benchmark.py b/torchvision/prototype/datasets/benchmark.py
@@ -105,7 +105,7 @@ def __init__(
         self.name = name
         self.variant = variant
 
-        self.new_raw_dataset = new_datasets._api.find(name)
+        self.new_raw_dataset = new_datasets._api._find(name)
         self.legacy_cls = legacy_cls or self._find_legacy_cls()
 
         if new_config is None:
diff --git a/torchvision/prototype/datasets/generate_category_files.py b/torchvision/prototype/datasets/generate_category_files.py
@@ -6,7 +6,7 @@
 import sys
 
 from torchvision.prototype import datasets
-from torchvision.prototype.datasets._api import find
+from torchvision.prototype.datasets._api import _find
 from torchvision.prototype.datasets.utils._internal import BUILTIN_DIR
 
 
@@ -18,7 +18,7 @@ def main(*names, force=False):
         if path.exists() and not force:
             continue
 
-        dataset = find(name)
+        dataset = _find(name)
         try:
             categories = dataset._generate_categories(home / name)
         except NotImplementedError:
diff --git a/torchvision/prototype/features/__init__.py b/torchvision/prototype/features/__init__.py
@@ -1,6 +1,6 @@
 from ._bounding_box import BoundingBox, BoundingBoxFormat
 from ._encoded import EncodedData, EncodedImage, EncodedVideo
-from ._feature import Feature
+from ._feature import _Feature
 from ._image import ColorSpace, Image
 from ._label import Label, OneHotLabel
 from ._segmentation_mask import SegmentationMask
diff --git a/torchvision/prototype/features/_bounding_box.py b/torchvision/prototype/features/_bounding_box.py
@@ -5,18 +5,16 @@
 import torch
 from torchvision.prototype.utils._internal import StrEnum
 
-from ._feature import Feature
+from ._feature import _Feature
 
 
 class BoundingBoxFormat(StrEnum):
-    # this is just for test purposes
-    _SENTINEL = -1
     XYXY = StrEnum.auto()
     XYWH = StrEnum.auto()
     CXCYWH = StrEnum.auto()
 
 
-class BoundingBox(Feature):
+class BoundingBox(_Feature):
     format: BoundingBoxFormat
     image_size: Tuple[int, int]
 
diff --git a/torchvision/prototype/features/_encoded.py b/torchvision/prototype/features/_encoded.py
@@ -6,13 +6,13 @@
 import torch
 from torchvision.prototype.utils._internal import fromfile, ReadOnlyTensorBuffer
 
-from ._feature import Feature
+from ._feature import _Feature
 from ._image import Image
 
 D = TypeVar("D", bound="EncodedData")
 
 
-class EncodedData(Feature):
+class EncodedData(_Feature):
     @classmethod
     def _to_tensor(cls, data: Any, *, dtype: Optional[torch.dtype], device: Optional[torch.device]) -> torch.Tensor:
         # TODO: warn / bail out if we encounter a tensor with shape other than (N,) or with dtype other than uint8?
diff --git a/torchvision/prototype/features/_feature.py b/torchvision/prototype/features/_feature.py
@@ -4,10 +4,10 @@
 from torch._C import _TensorBase, DisableTorchFunction
 
 
-F = TypeVar("F", bound="Feature")
+F = TypeVar("F", bound="_Feature")
 
 
-class Feature(torch.Tensor):
+class _Feature(torch.Tensor):
     _META_ATTRS: Set[str] = set()
     _metadata: Dict[str, Any]
 
@@ -17,7 +17,7 @@ def __init_subclass__(cls) -> None:
         By adding the metadata attributes as class annotations on subclasses of :class:`Feature`, this method adds
         properties to have the same convenient access as regular attributes.
 
-        >>> class Foo(Feature):
+        >>> class Foo(_Feature):
         ...     bar: str
         ...     baz: Optional[str]
         >>> foo = Foo()
@@ -28,10 +28,10 @@ def __init_subclass__(cls) -> None:
         """
         meta_attrs = {attr for attr in cls.__annotations__.keys() - cls.__dict__.keys() if not attr.startswith("_")}
         for super_cls in cls.__mro__[1:]:
-            if super_cls is Feature:
+            if super_cls is _Feature:
                 break
 
-            meta_attrs.update(cast(Type[Feature], super_cls)._META_ATTRS)
+            meta_attrs.update(cast(Type[_Feature], super_cls)._META_ATTRS)
 
         cls._META_ATTRS = meta_attrs
         for name in meta_attrs:
diff --git a/torchvision/prototype/features/_image.py b/torchvision/prototype/features/_image.py
@@ -10,18 +10,16 @@
 from torchvision.utils import make_grid
 
 from ._bounding_box import BoundingBox
-from ._feature import Feature
+from ._feature import _Feature
 
 
 class ColorSpace(StrEnum):
-    # this is just for test purposes
-    _SENTINEL = -1
     OTHER = 0
     GRAYSCALE = 1
     RGB = 3
 
 
-class Image(Feature):
+class Image(_Feature):
     color_space: ColorSpace
 
     def __new__(
diff --git a/torchvision/prototype/features/_label.py b/torchvision/prototype/features/_label.py
@@ -5,10 +5,10 @@
 import torch
 from torchvision.prototype.utils._internal import apply_recursively
 
-from ._feature import Feature
+from ._feature import _Feature
 
 
-class Label(Feature):
+class Label(_Feature):
     categories: Optional[Sequence[str]]
 
     def __new__(
@@ -37,7 +37,7 @@ def to_categories(self) -> Any:
         return apply_recursively(lambda idx: cast(Sequence[str], self.categories)[idx], self.tolist())
 
 
-class OneHotLabel(Feature):
+class OneHotLabel(_Feature):
     categories: Optional[Sequence[str]]
 
     def __new__(
diff --git a/torchvision/prototype/features/_segmentation_mask.py b/torchvision/prototype/features/_segmentation_mask.py
@@ -1,5 +1,5 @@
-from ._feature import Feature
+from ._feature import _Feature
 
 
-class SegmentationMask(Feature):
+class SegmentationMask(_Feature):
     pass
diff --git a/torchvision/prototype/transforms/functional/_augment.py b/torchvision/prototype/transforms/functional/_augment.py
@@ -7,7 +7,7 @@
 
 from ._utils import dispatch
 
-T = TypeVar("T", bound=features.Feature)
+T = TypeVar("T", bound=features._Feature)
 
 
 @dispatch(
diff --git a/torchvision/prototype/transforms/functional/_color.py b/torchvision/prototype/transforms/functional/_color.py
@@ -8,7 +8,7 @@
 
 from ._utils import dispatch
 
-T = TypeVar("T", bound=features.Feature)
+T = TypeVar("T", bound=features._Feature)
 
 
 @dispatch(
diff --git a/torchvision/prototype/transforms/functional/_geometry.py b/torchvision/prototype/transforms/functional/_geometry.py
@@ -8,7 +8,7 @@
 
 from ._utils import dispatch
 
-T = TypeVar("T", bound=features.Feature)
+T = TypeVar("T", bound=features._Feature)
 
 
 @dispatch(
diff --git a/torchvision/prototype/transforms/functional/_misc.py b/torchvision/prototype/transforms/functional/_misc.py
@@ -7,7 +7,7 @@
 
 from ._utils import dispatch
 
-T = TypeVar("T", bound=features.Feature)
+T = TypeVar("T", bound=features._Feature)
 
 
 @dispatch(
diff --git a/torchvision/prototype/transforms/functional/_utils.py b/torchvision/prototype/transforms/functional/_utils.py
@@ -6,7 +6,7 @@
 import torch.overrides
 from torchvision.prototype import features
 
-F = TypeVar("F", bound=features.Feature)
+F = TypeVar("F", bound=features._Feature)
 
 
 def dispatch(kernels: Dict[Any, Optional[Callable]]) -> Callable[[Callable[..., F]], Callable[..., F]]:
@@ -79,7 +79,7 @@ def inner_wrapper(input: F, *args: Any, **kwargs: Any) -> F:
             else:
                 output = kernel(input, *args, **kwargs)
 
-            if issubclass(feature_type, features.Feature) and type(output) is torch.Tensor:
+            if issubclass(feature_type, features._Feature) and type(output) is torch.Tensor:
                 output = feature_type.new_like(input, output)
 
             return output
diff --git a/torchvision/prototype/transforms/kernels/_geometry.py b/torchvision/prototype/transforms/kernels/_geometry.py
@@ -7,7 +7,7 @@
 from ._meta_conversion import convert_bounding_box_format
 
 
-T = TypeVar("T", bound=features.Feature)
+T = TypeVar("T", bound=features._Feature)
 
 
 horizontal_flip_image = _F.hflip

Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@`
`17`	`17`	`OnlineResource,`
`18`	`18`	`)`
`19`	`19`	`from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE, read_mat, hint_sharding, hint_shuffling`
`20`		`-from torchvision.prototype.features import Label, BoundingBox, Feature, EncodedImage`
	`20`	`+from torchvision.prototype.features import Label, BoundingBox, _Feature, EncodedImage`
`21`	`21`
`22`	`22`
`23`	`23`	`class Caltech101(Dataset):`
`@@ -95,7 +95,7 @@ def _prepare_sample(`
`95`	`95`	`bounding_box=BoundingBox(`
`96`	`96`	`ann["box_coord"].astype(np.int64).squeeze()[[2, 0, 3, 1]], format="xyxy", image_size=image.image_size`
`97`	`97`	`),`
`98`		`- contour=Feature(ann["obj_contour"].T),`
	`98`	`+ contour=_Feature(ann["obj_contour"].T),`
`99`	`99`	`)`
`100`	`100`
`101`	`101`	`def _make_datapipe(`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@`
`29`	`29`	`path_comparator,`
`30`	`30`	`path_accessor,`
`31`	`31`	`)`
`32`		`-from torchvision.prototype.features import Label, BoundingBox, Feature, EncodedImage`
	`32`	`+from torchvision.prototype.features import Label, BoundingBox, _Feature, EncodedImage`
`33`	`33`
`34`	`34`	`csv.register_dialect("cub200", delimiter=" ")`
`35`	`35`
`@@ -131,7 +131,7 @@ def _2010_prepare_ann(self, data: Tuple[str, Tuple[str, BinaryIO]], image_size:`
`131`	`131`	`format="xyxy",`
`132`	`132`	`image_size=image_size,`
`133`	`133`	`),`
`134`		`- segmentation=Feature(content["seg"]),`
	`134`	`+ segmentation=_Feature(content["seg"]),`
`135`	`135`	`)`
`136`	`136`
`137`	`137`	`def _prepare_sample(`