Skip to content

GitLab

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

Revamp prototype features and transforms (#5407) 

* revamp prototype features (#5283)

* remove decoding from prototype datasets (#5287)

* remove decoder from prototype datasets

* remove unused imports

* cleanup

* fix readme

* use OneHotLabel in SEMEION

* improve voc implementation

* revert unrelated changes

* fix semeion mock data

* fix pcam

* readd functional transforms API to prototype (#5295)

* readd functional transforms

* cleanup

* add missing imports

* remove __torch_function__ dispatch

* readd repr

* readd empty line

* add test for scriptability

* remove function copy

* change import from functional tensor transforms to just functional

* fix import

* fix test

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

* fix prototype functional transforms after review

* address features review

* make mypy more strict on prototype features

* make mypy more strict for prototype transforms

* fix annotation

* fix kernel tests

* add automatic feature type dispatch to functional transforms (#5323)

* add auto dispatch

* fix missing arguments error message

* remove pil kernel for erase

* automate feature specific parameter detection

* fix typos

* cleanup dispatcher call

* remove __torch_function__ from transform dispatch

* remove auto-generation

* revert unrelated changes

* remove implements decorator

* change register parameter order

* change order of transforms for readability

* add documentation for __torch_function__

* fix mypy

* inline check for support

* refactor kernel registering process

* refactor dispatch to be a regular decorator

* split kernels and dispatchers

* remove sentinels

* replace pass with ...

* appease mypy

* make single kernel dispatchers more concise

* make dispatcher signatures more generic

* make kernel checking more strict

* revert doc changes

* address Franciscos comments

* remove inplace

* rename kernel test module

* fix inplace

* remove special casing for pil and vanilla tensors

* address comments

* update docs

* cleanup features / transforms feature branch (#5406)

* mark candidates for removal

* align signature of resize_bounding_box with corresponding image kernel

* fix documentation of Feature

* remove interpolation mode and antialias option from resize_segmentation_mask

* remove or privatize functionality in features / datasets / transforms
parent f2f490b1
Hide whitespace changes
Inline Side-by-side
Showing with 443 additions and 597 deletions
torchvision/prototype/datasets/_builtin/mnist.py
View file @ d32bc4ba
import abc
import functools
import io
import operator
import pathlib
import string
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, cast, BinaryIO, Union, Sequence
from typing import Any, Dict, Iterator, List, Optional, Tuple, cast, BinaryIO, Union, Sequence
import torch
from torchdata.datapipes.iter import (
......@@ -13,24 +12,21 @@ from torchdata.datapipes.iter import (
Mapper,
Zipper,
)
from torchvision.prototype.datasets.decoder import raw
from torchvision.prototype.datasets.utils import (
Dataset,
DatasetType,
DatasetConfig,
DatasetInfo,
HttpResource,
OnlineResource,
)
from torchvision.prototype.datasets.utils._internal import (
image_buffer_from_array,
Decompressor,
INFINITE_BUFFER_SIZE,
fromfile,
hint_sharding,
hint_shuffling,
)
from torchvision.prototype.features import Image, Label
from torchvision.prototype.utils._internal import fromfile
__all__ = ["MNIST", "FashionMNIST", "KMNIST", "EMNIST", "QMNIST"]
......@@ -105,31 +101,15 @@ class _MNISTBase(Dataset):
def start_and_stop(self, config: DatasetConfig) -> Tuple[Optional[int], Optional[int]]:
return None, None
def _collate_and_decode(
self,
data: Tuple[torch.Tensor, torch.Tensor],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
def _prepare_sample(self, data: Tuple[torch.Tensor, torch.Tensor], *, config: DatasetConfig) -> Dict[str, Any]:
image, label = data
if decoder is raw:
image = Image(image)
else:
image_buffer = image_buffer_from_array(image.numpy())
image = decoder(image_buffer) if decoder else image_buffer # type: ignore[assignment]
label = Label(label, dtype=torch.int64, category=self.info.categories[int(label)])
return dict(image=image, label=label)
return dict(
image=Image(image),
label=Label(label, dtype=torch.int64, categories=self.categories),
)
def _make_datapipe(
self,
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
self, resource_dps: List[IterDataPipe], *, config: DatasetConfig
) -> IterDataPipe[Dict[str, Any]]:
images_dp, labels_dp = resource_dps
start, stop = self.start_and_stop(config)
......@@ -143,14 +123,13 @@ class _MNISTBase(Dataset):
dp = Zipper(images_dp, labels_dp)
dp = hint_sharding(dp)
dp = hint_shuffling(dp)
return Mapper(dp, functools.partial(self._collate_and_decode, config=config, decoder=decoder))
return Mapper(dp, functools.partial(self._prepare_sample, config=config))
class MNIST(_MNISTBase):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"mnist",
type=DatasetType.RAW,
categories=10,
homepage="http://yann.lecun.com/exdb/mnist",
valid_options=dict(
......@@ -183,7 +162,6 @@ class FashionMNIST(MNIST):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"fashionmnist",
type=DatasetType.RAW,
categories=(
"T-shirt/top",
"Trouser",
......@@ -215,7 +193,6 @@ class KMNIST(MNIST):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"kmnist",
type=DatasetType.RAW,
categories=["o", "ki", "su", "tsu", "na", "ha", "ma", "ya", "re", "wo"],
homepage="http://codh.rois.ac.jp/kmnist/index.html.en",
valid_options=dict(
......@@ -236,7 +213,6 @@ class EMNIST(_MNISTBase):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"emnist",
type=DatasetType.RAW,
categories=list(string.digits + string.ascii_uppercase + string.ascii_lowercase),
homepage="https://www.westernsydney.edu.au/icns/reproducible_research/publication_support_materials/emnist",
valid_options=dict(
......@@ -291,13 +267,7 @@ class EMNIST(_MNISTBase):
46: 9,
}
def _collate_and_decode(
self,
data: Tuple[torch.Tensor, torch.Tensor],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
def _prepare_sample(self, data: Tuple[torch.Tensor, torch.Tensor], *, config: DatasetConfig) -> Dict[str, Any]:
# In these two splits, some lowercase letters are merged into their uppercase ones (see Fig 2. in the paper).
# That means for example that there is 'D', 'd', and 'C', but not 'c'. Since the labels are nevertheless dense,
# i.e. no gaps between 0 and 46 for 47 total classes, we need to add an offset to create this gaps. For example,
......@@ -310,14 +280,10 @@ class EMNIST(_MNISTBase):
image, label = data
label += self._LABEL_OFFSETS.get(int(label), 0)
data = (image, label)
return super()._collate_and_decode(data, config=config, decoder=decoder)
return super()._prepare_sample(data, config=config)
def _make_datapipe(
self,
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
self, resource_dps: List[IterDataPipe], *, config: DatasetConfig
) -> IterDataPipe[Dict[str, Any]]:
archive_dp = resource_dps[0]
images_dp, labels_dp = Demultiplexer(
......@@ -327,14 +293,13 @@ class EMNIST(_MNISTBase):
drop_none=True,
buffer_size=INFINITE_BUFFER_SIZE,
)
return super()._make_datapipe([images_dp, labels_dp], config=config, decoder=decoder)
return super()._make_datapipe([images_dp, labels_dp], config=config)
class QMNIST(_MNISTBase):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"qmnist",
type=DatasetType.RAW,
categories=10,
homepage="https://github.com/facebookresearch/qmnist",
valid_options=dict(
......@@ -376,16 +341,10 @@ class QMNIST(_MNISTBase):
return start, stop
def _collate_and_decode(
self,
data: Tuple[torch.Tensor, torch.Tensor],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
def _prepare_sample(self, data: Tuple[torch.Tensor, torch.Tensor], *, config: DatasetConfig) -> Dict[str, Any]:
image, ann = data
label, *extra_anns = ann
sample = super()._collate_and_decode((image, label), config=config, decoder=decoder)
sample = super()._prepare_sample((image, label), config=config)
sample.update(
dict(
......
torchvision/prototype/datasets/_builtin/oxford_iiit_pet.py
View file @ d32bc4ba
import enum
import functools
import io
import pathlib
from typing import Any, Callable, Dict, List, Optional, Tuple
from typing import Any, Dict, List, Optional, Tuple, BinaryIO
import torch
from torchdata.datapipes.iter import IterDataPipe, Mapper, Filter, IterKeyZipper, Demultiplexer, CSVDictParser
from torchvision.prototype.datasets.utils import (
Dataset,
......@@ -12,7 +9,6 @@ from torchvision.prototype.datasets.utils import (
DatasetInfo,
HttpResource,
OnlineResource,
DatasetType,
)
from torchvision.prototype.datasets.utils._internal import (
INFINITE_BUFFER_SIZE,
......@@ -22,7 +18,7 @@ from torchvision.prototype.datasets.utils._internal import (
path_accessor,
path_comparator,
)
from torchvision.prototype.features import Label
from torchvision.prototype.features import Label, EncodedImage
class OxfordIITPetDemux(enum.IntEnum):
......@@ -34,7 +30,6 @@ class OxfordIITPet(Dataset):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"oxford-iiit-pet",
type=DatasetType.IMAGE,
homepage="https://www.robots.ox.ac.uk/~vgg/data/pets/",
valid_options=dict(
split=("trainval", "test"),
......@@ -66,18 +61,8 @@ class OxfordIITPet(Dataset):
def _filter_segmentations(self, data: Tuple[str, Any]) -> bool:
return not pathlib.Path(data[0]).name.startswith(".")
def _decode_classification_data(self, data: Dict[str, str]) -> Dict[str, Any]:
label_idx = int(data["label"]) - 1
return dict(
label=Label(label_idx, category=self.info.categories[label_idx]),
species="cat" if data["species"] == "1" else "dog",
)
def _collate_and_decode_sample(
self,
data: Tuple[Tuple[Dict[str, str], Tuple[str, io.IOBase]], Tuple[str, io.IOBase]],
*,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
def _prepare_sample(
self, data: Tuple[Tuple[Dict[str, str], Tuple[str, BinaryIO]], Tuple[str, BinaryIO]]
) -> Dict[str, Any]:
ann_data, image_data = data
classification_data, segmentation_data = ann_data
......@@ -85,19 +70,16 @@ class OxfordIITPet(Dataset):
image_path, image_buffer = image_data
return dict(
self._decode_classification_data(classification_data),
label=Label(int(classification_data["label"]) - 1, categories=self.categories),
species="cat" if classification_data["species"] == "1" else "dog",
segmentation_path=segmentation_path,
segmentation=decoder(segmentation_buffer) if decoder else segmentation_buffer,
segmentation=EncodedImage.from_file(segmentation_buffer),
image_path=image_path,
image=decoder(image_buffer) if decoder else image_buffer,
image=EncodedImage.from_file(image_buffer),
)
def _make_datapipe(
self,
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
self, resource_dps: List[IterDataPipe], *, config: DatasetConfig
) -> IterDataPipe[Dict[str, Any]]:
images_dp, anns_dp = resource_dps
......@@ -137,7 +119,7 @@ class OxfordIITPet(Dataset):
ref_key_fn=path_accessor("stem"),
buffer_size=INFINITE_BUFFER_SIZE,
)
return Mapper(dp, functools.partial(self._collate_and_decode_sample, decoder=decoder))
return Mapper(dp, self._prepare_sample)
def _filter_split_and_classification_anns(self, data: Tuple[str, Any]) -> bool:
return self._classify_anns(data) == OxfordIITPetDemux.SPLIT_AND_CLASSIFICATION
......
torchvision/prototype/datasets/_builtin/pcam.py
View file @ d32bc4ba
import io
from collections import namedtuple
from typing import Any, Callable, Dict, List, Optional, Tuple, Iterator
from typing import Any, Dict, List, Optional, Tuple, Iterator
import torch
from torchdata.datapipes.iter import IterDataPipe, Mapper, Zipper
from torchvision.prototype import features
from torchvision.prototype.datasets.utils import (
......@@ -10,7 +9,6 @@ from torchvision.prototype.datasets.utils import (
DatasetConfig,
DatasetInfo,
OnlineResource,
DatasetType,
GDriveResource,
)
from torchvision.prototype.datasets.utils._internal import (
......@@ -46,7 +44,6 @@ class PCAM(Dataset):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"pcam",
type=DatasetType.RAW,
homepage="https://github.com/basveeling/pcam",
categories=2,
valid_options=dict(split=("train", "test", "val")),
......@@ -98,7 +95,7 @@ class PCAM(Dataset):
for file_name, gdrive_id, sha256 in self._RESOURCES[config.split]
]
def _collate_and_decode(self, data: Tuple[Any, Any]) -> Dict[str, Any]:
def _prepare_sample(self, data: Tuple[Any, Any]) -> Dict[str, Any]:
image, target = data # They're both numpy arrays at this point
return {
......@@ -107,11 +104,7 @@ class PCAM(Dataset):
}
def _make_datapipe(
self,
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
self, resource_dps: List[IterDataPipe], *, config: DatasetConfig
) -> IterDataPipe[Dict[str, Any]]:
images_dp, targets_dp = resource_dps
......@@ -122,4 +115,4 @@ class PCAM(Dataset):
dp = Zipper(images_dp, targets_dp)
dp = hint_sharding(dp)
dp = hint_shuffling(dp)
return Mapper(dp, self._collate_and_decode)
return Mapper(dp, self._prepare_sample)
torchvision/prototype/datasets/_builtin/sbd.py
View file @ d32bc4ba
import functools
import io
import pathlib
import re
from typing import Any, Callable, Dict, List, Optional, Tuple, cast
from typing import Any, Dict, List, Optional, Tuple, cast, BinaryIO
import numpy as np
import torch
from torchdata.datapipes.iter import (
IterDataPipe,
Mapper,
......@@ -20,7 +17,6 @@ from torchvision.prototype.datasets.utils import (
DatasetInfo,
HttpResource,
OnlineResource,
DatasetType,
)
from torchvision.prototype.datasets.utils._internal import (
INFINITE_BUFFER_SIZE,
......@@ -31,20 +27,17 @@ from torchvision.prototype.datasets.utils._internal import (
hint_sharding,
hint_shuffling,
)
from torchvision.prototype.features import Feature
from torchvision.prototype.features import _Feature, EncodedImage
class SBD(Dataset):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"sbd",
type=DatasetType.IMAGE,
dependencies=("scipy",),
homepage="http://home.bharathh.info/pubs/codes/SBD/download.html",
valid_options=dict(
split=("train", "val", "train_noval"),
boundaries=(True, False),
segmentation=(False, True),
),
)
......@@ -75,50 +68,21 @@ class SBD(Dataset):
else:
return None
def _decode_ann(
self, data: Dict[str, Any], *, decode_boundaries: bool, decode_segmentation: bool
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
raw_anns = data["GTcls"][0]
raw_boundaries = raw_anns["Boundaries"][0]
raw_segmentation = raw_anns["Segmentation"][0]
# the boundaries are stored in sparse CSC format, which is not supported by PyTorch
boundaries = (
Feature(np.stack([raw_boundary[0].toarray() for raw_boundary in raw_boundaries]))
if decode_boundaries
else None
)
segmentation = Feature(raw_segmentation) if decode_segmentation else None
return boundaries, segmentation
def _collate_and_decode_sample(
self,
data: Tuple[Tuple[Any, Tuple[str, io.IOBase]], Tuple[str, io.IOBase]],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
def _prepare_sample(self, data: Tuple[Tuple[Any, Tuple[str, BinaryIO]], Tuple[str, BinaryIO]]) -> Dict[str, Any]:
split_and_image_data, ann_data = data
_, image_data = split_and_image_data
image_path, image_buffer = image_data
ann_path, ann_buffer = ann_data
image = decoder(image_buffer) if decoder else image_buffer
if config.boundaries or config.segmentation:
boundaries, segmentation = self._decode_ann(
read_mat(ann_buffer), decode_boundaries=config.boundaries, decode_segmentation=config.segmentation
)
else:
boundaries = segmentation = None
anns = read_mat(ann_buffer, squeeze_me=True)["GTcls"]
return dict(
image_path=image_path,
image=image,
image=EncodedImage.from_file(image_buffer),
ann_path=ann_path,
boundaries=boundaries,
segmentation=segmentation,
# 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()),
)
def _make_datapipe(
......@@ -126,7 +90,6 @@ class SBD(Dataset):
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> IterDataPipe[Dict[str, Any]]:
archive_dp, extra_split_dp = resource_dps
......@@ -138,10 +101,10 @@ class SBD(Dataset):
buffer_size=INFINITE_BUFFER_SIZE,
drop_none=True,
)
if config.split == "train_noval":
split_dp = extra_split_dp
split_dp = Filter(split_dp, path_comparator("stem", config.split))
split_dp = Filter(split_dp, path_comparator("name", f"{config.split}.txt"))
split_dp = LineReader(split_dp, decode=True)
split_dp = hint_sharding(split_dp)
split_dp = hint_shuffling(split_dp)
......@@ -155,7 +118,7 @@ class SBD(Dataset):
ref_key_fn=path_accessor("stem"),
buffer_size=INFINITE_BUFFER_SIZE,
)
return Mapper(dp, functools.partial(self._collate_and_decode_sample, config=config, decoder=decoder))
return Mapper(dp, self._prepare_sample)
def _generate_categories(self, root: pathlib.Path) -> Tuple[str, ...]:
resources = self.resources(self.default_config)
......
torchvision/prototype/datasets/_builtin/semeion.py
View file @ d32bc4ba
import functools
import io
from typing import Any, Callable, Dict, List, Optional, Tuple
from typing import Any, Dict, List, Tuple
import torch
from torchdata.datapipes.iter import (
......@@ -8,24 +6,21 @@ from torchdata.datapipes.iter import (
Mapper,
CSVParser,
)
from torchvision.prototype.datasets.decoder import raw
from torchvision.prototype.datasets.utils import (
Dataset,
DatasetConfig,
DatasetInfo,
HttpResource,
OnlineResource,
DatasetType,
)
from torchvision.prototype.datasets.utils._internal import image_buffer_from_array, hint_sharding, hint_shuffling
from torchvision.prototype.features import Image, Label
from torchvision.prototype.datasets.utils._internal import hint_sharding, hint_shuffling
from torchvision.prototype.features import Image, OneHotLabel
class SEMEION(Dataset):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"semeion",
type=DatasetType.RAW,
categories=10,
homepage="https://archive.ics.uci.edu/ml/datasets/Semeion+Handwritten+Digit",
)
......@@ -37,34 +32,22 @@ class SEMEION(Dataset):
)
return [data]
def _collate_and_decode_sample(
self,
data: Tuple[str, ...],
*,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
image_data = torch.tensor([float(pixel) for pixel in data[:256]], dtype=torch.uint8).reshape(16, 16)
label_data = [int(label) for label in data[256:] if label]
def _prepare_sample(self, data: Tuple[str, ...]) -> Dict[str, Any]:
image_data, label_data = data[:256], data[256:-1]
if decoder is raw:
image = Image(image_data.unsqueeze(0))
else:
image_buffer = image_buffer_from_array(image_data.numpy())
image = decoder(image_buffer) if decoder else image_buffer # type: ignore[assignment]
label_idx = next((idx for idx, one_hot_label in enumerate(label_data) if one_hot_label))
return dict(image=image, label=Label(label_idx, category=self.info.categories[label_idx]))
return dict(
image=Image(torch.tensor([float(pixel) for pixel in image_data], dtype=torch.uint8).reshape(16, 16)),
label=OneHotLabel([int(label) for label in label_data], categories=self.categories),
)
def _make_datapipe(
self,
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> IterDataPipe[Dict[str, Any]]:
dp = resource_dps[0]
dp = CSVParser(dp, delimiter=" ")
dp = hint_sharding(dp)
dp = hint_shuffling(dp)
dp = Mapper(dp, functools.partial(self._collate_and_decode_sample, decoder=decoder))
return dp
return Mapper(dp, self._prepare_sample)
torchvision/prototype/datasets/_builtin/svhn.py
View file @ d32bc4ba
import functools
import io
from typing import Any, Callable, Dict, List, Optional, Tuple
from typing import Any, Dict, List, Tuple, BinaryIO
import numpy as np
import torch
from torchdata.datapipes.iter import (
IterDataPipe,
Mapper,
UnBatcher,
)
from torchvision.prototype.datasets.decoder import raw
from torchvision.prototype.datasets.utils import (
Dataset,
DatasetConfig,
DatasetInfo,
HttpResource,
OnlineResource,
DatasetType,
)
from torchvision.prototype.datasets.utils._internal import (
read_mat,
hint_sharding,
hint_shuffling,
image_buffer_from_array,
)
from torchvision.prototype.features import Label, Image
......@@ -31,7 +25,6 @@ class SVHN(Dataset):
def _make_info(self) -> DatasetInfo:
return DatasetInfo(
"svhn",
type=DatasetType.RAW,
dependencies=("scipy",),
categories=10,
homepage="http://ufldl.stanford.edu/housenumbers/",
......@@ -52,7 +45,7 @@ class SVHN(Dataset):
return [data]
def _read_images_and_labels(self, data: Tuple[str, io.IOBase]) -> List[Tuple[np.ndarray, np.ndarray]]:
def _read_images_and_labels(self, data: Tuple[str, BinaryIO]) -> List[Tuple[np.ndarray, np.ndarray]]:
_, buffer = data
content = read_mat(buffer)
return list(
......@@ -62,23 +55,12 @@ class SVHN(Dataset):
)
)
def _collate_and_decode_sample(
self,
data: Tuple[np.ndarray, np.ndarray],
*,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
def _prepare_sample(self, data: Tuple[np.ndarray, np.ndarray]) -> Dict[str, Any]:
image_array, label_array = data
if decoder is raw:
image = Image(image_array.transpose((2, 0, 1)))
else:
image_buffer = image_buffer_from_array(image_array)
image = decoder(image_buffer) if decoder else image_buffer # type: ignore[assignment]
return dict(
image=image,
label=Label(int(label_array) % 10),
image=Image(image_array.transpose((2, 0, 1))),
label=Label(int(label_array) % 10, categories=self.categories),
)
def _make_datapipe(
......@@ -86,11 +68,10 @@ class SVHN(Dataset):
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> IterDataPipe[Dict[str, Any]]:
dp = resource_dps[0]
dp = Mapper(dp, self._read_images_and_labels)
dp = UnBatcher(dp)
dp = hint_sharding(dp)
dp = hint_shuffling(dp)
return Mapper(dp, functools.partial(self._collate_and_decode_sample, decoder=decoder))
return Mapper(dp, self._prepare_sample)
torchvision/prototype/datasets/_builtin/voc.categories 0 → 100644
View file @ d32bc4ba
aeroplane
bicycle
bird
boat
bottle
bus
car
cat
chair
cow
diningtable
dog
horse
motorbike
person
pottedplant
sheep
sofa
train
tvmonitor
torchvision/prototype/datasets/_builtin/voc.py
View file @ d32bc4ba
import functools
import io
import pathlib
from typing import Any, Callable, Dict, List, Optional, Tuple
from typing import Any, Dict, List, Optional, Tuple, BinaryIO, cast, Callable
from xml.etree import ElementTree
import torch
from torchdata.datapipes.iter import (
IterDataPipe,
Mapper,
......@@ -20,7 +18,6 @@ from torchvision.prototype.datasets.utils import (
DatasetInfo,
HttpResource,
OnlineResource,
DatasetType,
)
from torchvision.prototype.datasets.utils._internal import (
path_accessor,
......@@ -30,7 +27,7 @@ from torchvision.prototype.datasets.utils._internal import (
hint_sharding,
hint_shuffling,
)
from torchvision.prototype.features import BoundingBox
from torchvision.prototype.features import BoundingBox, Label, EncodedImage
class VOCDatasetInfo(DatasetInfo):
......@@ -50,7 +47,6 @@ class VOC(Dataset):
def _make_info(self) -> DatasetInfo:
return VOCDatasetInfo(
"voc",
type=DatasetType.IMAGE,
homepage="http://host.robots.ox.ac.uk/pascal/VOC/",
valid_options=dict(
split=("train", "val", "trainval", "test"),
......@@ -99,40 +95,52 @@ class VOC(Dataset):
else:
return None
def _decode_detection_ann(self, buffer: io.IOBase) -> torch.Tensor:
result = VOCDetection.parse_voc_xml(ElementTree.parse(buffer).getroot()) # type: ignore[arg-type]
objects = result["annotation"]["object"]
bboxes = [obj["bndbox"] for obj in objects]
bboxes = [[int(bbox[part]) for part in ("xmin", "ymin", "xmax", "ymax")] for bbox in bboxes]
return BoundingBox(bboxes)
def _parse_detection_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
return cast(Dict[str, Any], VOCDetection.parse_voc_xml(ElementTree.parse(buffer).getroot())["annotation"])
def _prepare_detection_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
anns = self._parse_detection_ann(buffer)
instances = anns["object"]
return dict(
bounding_boxes=BoundingBox(
[
[int(instance["bndbox"][part]) for part in ("xmin", "ymin", "xmax", "ymax")]
for instance in instances
],
format="xyxy",
image_size=cast(Tuple[int, int], tuple(int(anns["size"][dim]) for dim in ("height", "width"))),
),
labels=Label(
[self.categories.index(instance["name"]) for instance in instances], categories=self.categories
),
)
def _prepare_segmentation_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
return dict(segmentation=EncodedImage.from_file(buffer))
def _collate_and_decode_sample(
def _prepare_sample(
self,
data: Tuple[Tuple[Tuple[str, str], Tuple[str, io.IOBase]], Tuple[str, io.IOBase]],
data: Tuple[Tuple[Tuple[str, str], Tuple[str, BinaryIO]], Tuple[str, BinaryIO]],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
prepare_ann_fn: Callable[[BinaryIO], Dict[str, Any]],
) -> Dict[str, Any]:
split_and_image_data, ann_data = data
_, image_data = split_and_image_data
image_path, image_buffer = image_data
ann_path, ann_buffer = ann_data
image = decoder(image_buffer) if decoder else image_buffer
if config.task == "detection":
ann = self._decode_detection_ann(ann_buffer)
else: # config.task == "segmentation":
ann = decoder(ann_buffer) if decoder else ann_buffer # type: ignore[assignment]
return dict(image_path=image_path, image=image, ann_path=ann_path, ann=ann)
return dict(
prepare_ann_fn(ann_buffer),
image_path=image_path,
image=EncodedImage.from_file(image_buffer),
ann_path=ann_path,
)
def _make_datapipe(
self,
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> IterDataPipe[Dict[str, Any]]:
archive_dp = resource_dps[0]
split_dp, images_dp, anns_dp = Demultiplexer(
......@@ -158,4 +166,25 @@ class VOC(Dataset):
ref_key_fn=path_accessor("stem"),
buffer_size=INFINITE_BUFFER_SIZE,
)
return Mapper(dp, functools.partial(self._collate_and_decode_sample, config=config, decoder=decoder))
return Mapper(
dp,
functools.partial(
self._prepare_sample,
prepare_ann_fn=self._prepare_detection_ann
if config.task == "detection"
else self._prepare_segmentation_ann,
),
)
def _filter_detection_anns(self, data: Tuple[str, Any], *, config: DatasetConfig) -> bool:
return self._classify_archive(data, config=config) == 2
def _generate_categories(self, root: pathlib.Path) -> List[str]:
config = self.info.make_config(task="detection")
resource = self.resources(config)[0]
dp = resource.load(pathlib.Path(root) / self.name)
dp = Filter(dp, self._filter_detection_anns, fn_kwargs=dict(config=config))
dp = Mapper(dp, self._parse_detection_ann, input_col=1)
return sorted({instance["name"] for _, anns in dp for instance in anns["object"]})
torchvision/prototype/datasets/_folder.py
View file @ d32bc4ba
import functools
import io
import os
import os.path
import pathlib
from typing import Callable, Optional, Collection
from typing import Union, Tuple, List, Dict, Any
from typing import BinaryIO, Optional, Collection, Union, Tuple, List, Dict, Any
import torch
from torchdata.datapipes.iter import IterDataPipe, FileLister, FileOpener, Mapper, Shuffler, Filter
from torchvision.prototype.datasets.decoder import pil
from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE, hint_sharding
from torchdata.datapipes.iter import IterDataPipe, FileLister, Mapper, Filter, FileOpener
from torchvision.prototype.datasets.utils._internal import hint_sharding, hint_shuffling
from torchvision.prototype.features import Label, EncodedImage, EncodedData
__all__ = ["from_data_folder", "from_image_folder"]
......@@ -20,29 +17,24 @@ def _is_not_top_level_file(path: str, *, root: pathlib.Path) -> bool:
return rel_path.is_dir() or rel_path.parent != pathlib.Path(".")
def _collate_and_decode_data(
data: Tuple[str, io.IOBase],
def _prepare_sample(
data: Tuple[str, BinaryIO],
*,
root: pathlib.Path,
categories: List[str],
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> Dict[str, Any]:
path, buffer = data
data = decoder(buffer) if decoder else buffer
category = pathlib.Path(path).relative_to(root).parts[0]
label = torch.tensor(categories.index(category))
return dict(
path=path,
data=data,
label=label,
category=category,
data=EncodedData.from_file(buffer),
label=Label.from_category(category, categories=categories),
)
def from_data_folder(
root: Union[str, pathlib.Path],
*,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]] = None,
valid_extensions: Optional[Collection[str]] = None,
recursive: bool = True,
) -> Tuple[IterDataPipe, List[str]]:
......@@ -52,26 +44,22 @@ def from_data_folder(
dp = FileLister(str(root), recursive=recursive, masks=masks)
dp: IterDataPipe = Filter(dp, functools.partial(_is_not_top_level_file, root=root))
dp = hint_sharding(dp)
dp = Shuffler(dp, buffer_size=INFINITE_BUFFER_SIZE)
dp = hint_shuffling(dp)
dp = FileOpener(dp, mode="rb")
return (
Mapper(dp, functools.partial(_collate_and_decode_data, root=root, categories=categories, decoder=decoder)),
categories,
)
return Mapper(dp, functools.partial(_prepare_sample, root=root, categories=categories)), categories
def _data_to_image_key(sample: Dict[str, Any]) -> Dict[str, Any]:
sample["image"] = sample.pop("data")
sample["image"] = EncodedImage(sample.pop("data").data)
return sample
def from_image_folder(
root: Union[str, pathlib.Path],
*,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]] = pil,
valid_extensions: Collection[str] = ("jpg", "jpeg", "png", "ppm", "bmp", "pgm", "tif", "tiff", "webp"),
**kwargs: Any,
) -> Tuple[IterDataPipe, List[str]]:
valid_extensions = [valid_extension for ext in valid_extensions for valid_extension in (ext.lower(), ext.upper())]
dp, categories = from_data_folder(root, decoder=decoder, valid_extensions=valid_extensions, **kwargs)
dp, categories = from_data_folder(root, valid_extensions=valid_extensions, **kwargs)
return Mapper(dp, _data_to_image_key), categories
torchvision/prototype/datasets/decoder.py deleted 100644 → 0
View file @ f2f490b1
import io
import PIL.Image
import torch
from torchvision.prototype import features
from torchvision.transforms.functional import pil_to_tensor
__all__ = ["raw", "pil"]
def raw(buffer: io.IOBase) -> torch.Tensor:
raise RuntimeError("This is just a sentinel and should never be called.")
def pil(buffer: io.IOBase) -> features.Image:
return features.Image(pil_to_tensor(PIL.Image.open(buffer)))
torchvision/prototype/datasets/utils/__init__.py
View file @ d32bc4ba
from . import _internal
from ._dataset import DatasetType, DatasetConfig, DatasetInfo, Dataset
from . import _internal # usort: skip
from ._dataset import DatasetConfig, DatasetInfo, Dataset
from ._query import SampleQuery
from ._resource import OnlineResource, HttpResource, GDriveResource, ManualDownloadResource, KaggleDownloadResource
torchvision/prototype/datasets/utils/_dataset.py
View file @ d32bc4ba
import abc
import csv
import enum
import importlib
import io
import itertools
import os
import pathlib
from typing import Any, Callable, Dict, List, Optional, Sequence, Union, Tuple, Collection
from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Collection
import torch
from torch.utils.data import IterDataPipe
from torchvision.prototype.utils._internal import FrozenBunch, make_repr
from torchvision.prototype.utils._internal import add_suggestion, sequence_to_str
from torchvision.prototype.utils._internal import FrozenBunch, make_repr, add_suggestion, sequence_to_str
from .._home import use_sharded_dataset
from ._internal import BUILTIN_DIR, _make_sharded_datapipe
from ._resource import OnlineResource
class DatasetType(enum.Enum):
RAW = enum.auto()
IMAGE = enum.auto()
class DatasetConfig(FrozenBunch):
# This needs to be Frozen because we often pass configs as partial(func, config=config)
# and partial() requires the parameters to be hashable.
......@@ -34,7 +25,6 @@ class DatasetInfo:
self,
name: str,
*,
type: Union[str, DatasetType],
dependencies: Collection[str] = (),
categories: Optional[Union[int, Sequence[str], str, pathlib.Path]] = None,
citation: Optional[str] = None,
......@@ -44,7 +34,6 @@ class DatasetInfo:
extra: Optional[Dict[str, Any]] = None,
) -> None:
self.name = name.lower()
self.type = DatasetType[type.upper()] if isinstance(type, str) else type
self.dependecies = dependencies
......@@ -163,7 +152,6 @@ class Dataset(abc.ABC):
resource_dps: List[IterDataPipe],
*,
config: DatasetConfig,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
) -> IterDataPipe[Dict[str, Any]]:
pass
......@@ -175,7 +163,6 @@ class Dataset(abc.ABC):
root: Union[str, pathlib.Path],
*,
config: Optional[DatasetConfig] = None,
decoder: Optional[Callable[[io.IOBase], torch.Tensor]] = None,
skip_integrity_check: bool = False,
) -> IterDataPipe[Dict[str, Any]]:
if not config:
......@@ -190,7 +177,7 @@ class Dataset(abc.ABC):
resource_dps = [
resource.load(root, skip_integrity_check=skip_integrity_check) for resource in self.resources(config)
]
return self._make_datapipe(resource_dps, config=config, decoder=decoder)
return self._make_datapipe(resource_dps, config=config)
def _generate_categories(self, root: pathlib.Path) -> Sequence[Union[str, Sequence[str]]]:
raise NotImplementedError
torchvision/prototype/datasets/utils/_internal.py
View file @ d32bc4ba
import enum
import functools
import gzip
import io
import lzma
import mmap
import os
import os.path
import pathlib
import pickle
import platform
from typing import BinaryIO
from typing import (
Sequence,
......@@ -25,27 +22,24 @@ from typing import (
)
from typing import cast
import numpy as np
import PIL.Image
import torch
import torch.distributed as dist
import torch.utils.data
from torchdata.datapipes.iter import IoPathFileLister, IoPathFileOpener, IterDataPipe, ShardingFilter, Shuffler
from torchdata.datapipes.utils import StreamWrapper
from torchvision.prototype.utils._internal import fromfile
__all__ = [
"INFINITE_BUFFER_SIZE",
"BUILTIN_DIR",
"read_mat",
"image_buffer_from_array",
"MappingIterator",
"Enumerator",
"getitem",
"path_accessor",
"path_comparator",
"Decompressor",
"fromfile",
"read_flo",
"hint_sharding",
]
......@@ -59,7 +53,7 @@ INFINITE_BUFFER_SIZE = 1_000_000_000
BUILTIN_DIR = pathlib.Path(__file__).parent.parent / "_builtin"
def read_mat(buffer: io.IOBase, **kwargs: Any) -> Any:
def read_mat(buffer: BinaryIO, **kwargs: Any) -> Any:
try:
import scipy.io as sio
except ImportError as error:
......@@ -71,14 +65,6 @@ def read_mat(buffer: io.IOBase, **kwargs: Any) -> Any:
return sio.loadmat(buffer, **kwargs)
def image_buffer_from_array(array: np.ndarray, *, format: str = "png") -> io.BytesIO:
image = PIL.Image.fromarray(array)
buffer = io.BytesIO()
image.save(buffer, format=format)
buffer.seek(0)
return buffer
class MappingIterator(IterDataPipe[Union[Tuple[K, D], D]]):
def __init__(self, datapipe: IterDataPipe[Dict[K, D]], *, drop_key: bool = False) -> None:
self.datapipe = datapipe
......@@ -142,17 +128,17 @@ class CompressionType(enum.Enum):
LZMA = "lzma"
class Decompressor(IterDataPipe[Tuple[str, io.IOBase]]):
class Decompressor(IterDataPipe[Tuple[str, BinaryIO]]):
types = CompressionType
_DECOMPRESSORS = {
types.GZIP: lambda file: gzip.GzipFile(fileobj=file),
types.LZMA: lambda file: lzma.LZMAFile(file),
_DECOMPRESSORS: Dict[CompressionType, Callable[[BinaryIO], BinaryIO]] = {
types.GZIP: lambda file: cast(BinaryIO, gzip.GzipFile(fileobj=file)),
types.LZMA: lambda file: cast(BinaryIO, lzma.LZMAFile(file)),
}
def __init__(
self,
datapipe: IterDataPipe[Tuple[str, io.IOBase]],
datapipe: IterDataPipe[Tuple[str, BinaryIO]],
*,
type: Optional[Union[str, CompressionType]] = None,
) -> None:
......@@ -174,7 +160,7 @@ class Decompressor(IterDataPipe[Tuple[str, io.IOBase]]):
else:
raise RuntimeError("FIXME")
def __iter__(self) -> Iterator[Tuple[str, io.IOBase]]:
def __iter__(self) -> Iterator[Tuple[str, BinaryIO]]:
for path, file in self.datapipe:
type = self._detect_compression_type(path)
decompressor = self._DECOMPRESSORS[type]
......@@ -257,69 +243,6 @@ def _make_sharded_datapipe(root: str, dataset_size: int) -> IterDataPipe[Dict[st
return dp
def _read_mutable_buffer_fallback(file: BinaryIO, count: int, item_size: int) -> bytearray:
# A plain file.read() will give a read-only bytes, so we convert it to bytearray to make it mutable
return bytearray(file.read(-1 if count == -1 else count * item_size))
def fromfile(
file: BinaryIO,
*,
dtype: torch.dtype,
byte_order: str,
count: int = -1,
) -> torch.Tensor:
"""Construct a tensor from a binary file.
.. note::
This function is similar to :func:`numpy.fromfile` with two notable differences:
1. This function only accepts an open binary file, but not a path to it.
2. This function has an additional ``byte_order`` parameter, since PyTorch's ``dtype``'s do not support that
concept.
.. note::
If the ``file`` was opened in update mode, i.e. "r+b" or "w+b", reading data is much faster. Be aware that as
long as the file is still open, inplace operations on the returned tensor will reflect back to the file.
Args:
file (IO): Open binary file.
dtype (torch.dtype): Data type of the underlying data as well as of the returned tensor.
byte_order (str): Byte order of the data. Can be "little" or "big" endian.
count (int): Number of values of the returned tensor. If ``-1`` (default), will read the complete file.
"""
byte_order = "<" if byte_order == "little" else ">"
char = "f" if dtype.is_floating_point else ("i" if dtype.is_signed else "u")
item_size = (torch.finfo if dtype.is_floating_point else torch.iinfo)(dtype).bits // 8
np_dtype = byte_order + char + str(item_size)
buffer: Union[memoryview, bytearray]
if platform.system() != "Windows":
# PyTorch does not support tensors with underlying read-only memory. In case
# - the file has a .fileno(),
# - the file was opened for updating, i.e. 'r+b' or 'w+b',
# - the file is seekable
# we can avoid copying the data for performance. Otherwise we fall back to simply .read() the data and copy it
# to a mutable location afterwards.
try:
buffer = memoryview(mmap.mmap(file.fileno(), 0))[file.tell() :]
# Reading from the memoryview does not advance the file cursor, so we have to do it manually.
file.seek(*(0, io.SEEK_END) if count == -1 else (count * item_size, io.SEEK_CUR))
except (PermissionError, io.UnsupportedOperation):
buffer = _read_mutable_buffer_fallback(file, count, item_size)
else:
# On Windows just trying to call mmap.mmap() on a file that does not support it, may corrupt the internal state
# so no data can be read afterwards. Thus, we simply ignore the possible speed-up.
buffer = _read_mutable_buffer_fallback(file, count, item_size)
# We cannot use torch.frombuffer() directly, since it only supports the native byte order of the system. Thus, we
# read the data with np.frombuffer() with the correct byte order and convert it to the native one with the
# successive .astype() call.
return torch.from_numpy(np.frombuffer(buffer, dtype=np_dtype, count=count).astype(np_dtype[1:], copy=False))
def read_flo(file: BinaryIO) -> torch.Tensor:
if file.read(4) != b"PIEH":
raise ValueError("Magic number incorrect. Invalid .flo file")
......@@ -329,9 +252,9 @@ def read_flo(file: BinaryIO) -> torch.Tensor:
return flow.reshape((height, width, 2)).permute((2, 0, 1))
def hint_sharding(datapipe: IterDataPipe[D]) -> IterDataPipe[D]:
def hint_sharding(datapipe: IterDataPipe) -> ShardingFilter:
return ShardingFilter(datapipe)
def hint_shuffling(datapipe: IterDataPipe[D]) -> IterDataPipe[D]:
def hint_shuffling(datapipe: IterDataPipe[D]) -> Shuffler[D]:
return Shuffler(datapipe, default=False, buffer_size=INFINITE_BUFFER_SIZE)
torchvision/prototype/features/__init__.py
View file @ d32bc4ba
from ._bounding_box import BoundingBoxFormat, BoundingBox
from ._feature import Feature, DEFAULT
from ._image import Image, ColorSpace
from ._label import Label
from ._bounding_box import BoundingBox, BoundingBoxFormat
from ._encoded import EncodedData, EncodedImage, EncodedVideo
from ._feature import _Feature
from ._image import ColorSpace, Image
from ._label import Label, OneHotLabel
from ._segmentation_mask import SegmentationMask
torchvision/prototype/features/_bounding_box.py
View file @ d32bc4ba
import enum
import functools
from typing import Callable, Union, Tuple, Dict, Any, Optional, cast
from __future__ import annotations
from typing import Any, Tuple, Union, Optional
import torch
from torchvision.prototype.utils._internal import StrEnum
from ._feature import Feature, DEFAULT
from ._feature import _Feature
class BoundingBoxFormat(StrEnum):
# this is just for test purposes
_SENTINEL = -1
XYXY = enum.auto()
XYWH = enum.auto()
CXCYWH = enum.auto()
def to_parts(input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
return input.unbind(dim=-1) # type: ignore[return-value]
def from_parts(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
return torch.stack((a, b, c, d), dim=-1)
def format_converter_wrapper(
part_converter: Callable[
[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
]
):
def wrapper(input: torch.Tensor) -> torch.Tensor:
return from_parts(*part_converter(*to_parts(input)))
return wrapper
@format_converter_wrapper
def xywh_to_xyxy(
x: torch.Tensor, y: torch.Tensor, w: torch.Tensor, h: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
x1 = x
y1 = y
x2 = x + w
y2 = y + h
return x1, y1, x2, y2
@format_converter_wrapper
def xyxy_to_xywh(
x1: torch.Tensor, y1: torch.Tensor, x2: torch.Tensor, y2: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
x = x1
y = y1
w = x2 - x1
h = y2 - y1
return x, y, w, h
XYXY = StrEnum.auto()
XYWH = StrEnum.auto()
CXCYWH = StrEnum.auto()
@format_converter_wrapper
def cxcywh_to_xyxy(
cx: torch.Tensor, cy: torch.Tensor, w: torch.Tensor, h: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
x1 = cx - 0.5 * w
y1 = cy - 0.5 * h
x2 = cx + 0.5 * w
y2 = cy + 0.5 * h
return x1, y1, x2, y2
@format_converter_wrapper
def xyxy_to_cxcywh(
x1: torch.Tensor, y1: torch.Tensor, x2: torch.Tensor, y2: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
cx = (x1 + x2) / 2
cy = (y1 + y2) / 2
w = x2 - x1
h = y2 - y1
return cx, cy, w, h
class BoundingBox(Feature):
formats = BoundingBoxFormat
class BoundingBox(_Feature):
format: BoundingBoxFormat
image_size: Tuple[int, int]
@classmethod
def _parse_meta_data(
def __new__(
cls,
format: Union[str, BoundingBoxFormat] = DEFAULT, # type: ignore[assignment]
image_size: Optional[Tuple[int, int]] = DEFAULT, # type: ignore[assignment]
) -> Dict[str, Tuple[Any, Any]]:
data: Any,
*,
dtype: Optional[torch.dtype] = None,
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):
format = BoundingBoxFormat[format]
format_fallback = BoundingBoxFormat.XYXY
return dict(
format=(format, format_fallback),
image_size=(image_size, functools.partial(cls.guess_image_size, format=format_fallback)),
)
_TO_XYXY_MAP = {
BoundingBoxFormat.XYWH: xywh_to_xyxy,
BoundingBoxFormat.CXCYWH: cxcywh_to_xyxy,
}
_FROM_XYXY_MAP = {
BoundingBoxFormat.XYWH: xyxy_to_xywh,
BoundingBoxFormat.CXCYWH: xyxy_to_cxcywh,
}
bounding_box._metadata.update(dict(format=format, image_size=image_size))
@classmethod
def guess_image_size(cls, data: torch.Tensor, *, format: BoundingBoxFormat) -> Tuple[int, int]:
if format not in (BoundingBoxFormat.XYWH, BoundingBoxFormat.CXCYWH):
if format != BoundingBoxFormat.XYXY:
data = cls._TO_XYXY_MAP[format](data)
data = cls._FROM_XYXY_MAP[BoundingBoxFormat.XYWH](data)
*_, w, h = to_parts(data)
if data.dtype.is_floating_point:
w = w.ceil()
h = h.ceil()
return int(h.max()), int(w.max())
return bounding_box
@classmethod
def from_parts(
cls,
a,
b,
c,
d,
*,
like: Optional["BoundingBox"] = None,
format: Union[str, BoundingBoxFormat] = DEFAULT, # type: ignore[assignment]
image_size: Optional[Tuple[int, int]] = DEFAULT, # type: ignore[assignment]
) -> "BoundingBox":
return cls(from_parts(a, b, c, d), like=like, image_size=image_size, format=format)
def to_format(self, format: Union[str, BoundingBoxFormat]) -> BoundingBox:
# TODO: this is useful for developing and debugging but we should remove or at least revisit this before we
# promote this out of the prototype state
def to_parts(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
return to_parts(self)
# import at runtime to avoid cyclic imports
from torchvision.prototype.transforms.kernels import convert_bounding_box_format
def convert(self, format: Union[str, BoundingBoxFormat]) -> "BoundingBox":
if isinstance(format, str):
format = BoundingBoxFormat[format]
if format == self.format:
return cast(BoundingBox, self.clone())
data = self
if self.format != BoundingBoxFormat.XYXY:
data = self._TO_XYXY_MAP[self.format](data)
if format != BoundingBoxFormat.XYXY:
data = self._FROM_XYXY_MAP[format](data)
return BoundingBox(data, like=self, format=format)
return BoundingBox.new_like(
self, convert_bounding_box_format(self, old_format=self.format, new_format=format), format=format
)
torchvision/prototype/features/_encoded.py 0 → 100644
View file @ d32bc4ba
import os
import sys
from typing import BinaryIO, Tuple, Type, TypeVar, Union, Optional, Any
import PIL.Image
import torch
from torchvision.prototype.utils._internal import fromfile, ReadOnlyTensorBuffer
from ._feature import _Feature
from ._image import Image
D = TypeVar("D", bound="EncodedData")
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?
return super()._to_tensor(data, dtype=dtype, device=device)
@classmethod
def from_file(cls: Type[D], file: BinaryIO) -> D:
return cls(fromfile(file, dtype=torch.uint8, byte_order=sys.byteorder))
@classmethod
def from_path(cls: Type[D], path: Union[str, os.PathLike]) -> D:
with open(path, "rb") as file:
return cls.from_file(file)
class EncodedImage(EncodedData):
# TODO: Use @functools.cached_property if we can depend on Python 3.8
@property
def image_size(self) -> Tuple[int, int]:
if not hasattr(self, "_image_size"):
with PIL.Image.open(ReadOnlyTensorBuffer(self)) as image:
self._image_size = image.height, image.width
return self._image_size
def decode(self) -> Image:
# TODO: this is useful for developing and debugging but we should remove or at least revisit this before we
# promote this out of the prototype state
# import at runtime to avoid cyclic imports
from torchvision.prototype.transforms.kernels import decode_image_with_pil
return Image(decode_image_with_pil(self))
class EncodedVideo(EncodedData):
pass
torchvision/prototype/features/_feature.py
View file @ d32bc4ba
from typing import Tuple, cast, TypeVar, Set, Dict, Any, Callable, Optional, Mapping, Type, Sequence
from typing import Any, cast, Dict, Set, TypeVar, Union, Optional, Type, Callable, Tuple, Sequence, Mapping
import torch
from torch._C import _TensorBase, DisableTorchFunction
from torchvision.prototype.utils._internal import add_suggestion
F = TypeVar("F", bound="Feature")
F = TypeVar("F", bound="_Feature")
DEFAULT = object()
class Feature(torch.Tensor):
class _Feature(torch.Tensor):
_META_ATTRS: Set[str] = set()
_meta_data: Dict[str, Any]
def __init_subclass__(cls):
# In order to help static type checkers, we require subclasses of `Feature` add the meta data attributes
# as static class annotations:
#
# >>> class Foo(Feature):
# ... bar: str
# ... baz: Optional[str]
#
# Internally, this information is used twofold:
#
# 1. A class annotation is contained in `cls.__annotations__` but not in `cls.__dict__`. We use this difference
# to automatically detect the meta data attributes and expose them as `@property`'s for convenient runtime
# access. This happens in this method.
# 2. The information extracted in 1. is also used at creation (`__new__`) to perform an input parsing for
# unknown arguments.
_metadata: Dict[str, Any]
def __init_subclass__(cls) -> None:
"""
For convenient copying of metadata, we store it inside a dictionary rather than multiple individual attributes.
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):
... bar: str
... baz: Optional[str]
>>> foo = Foo()
>>> foo.bar
>>> foo.baz
This has the additional benefit that autocomplete engines and static type checkers are aware of the metadata.
"""
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(super_cls._META_ATTRS)
meta_attrs.update(cast(Type[_Feature], super_cls)._META_ATTRS)
cls._META_ATTRS = meta_attrs
for attr in meta_attrs:
setattr(cls, attr, property(lambda self, attr=attr: self._meta_data[attr]))
def __new__(cls, data, *, dtype=None, device=None, like=None, **kwargs):
unknown_meta_attrs = kwargs.keys() - cls._META_ATTRS
if unknown_meta_attrs:
unknown_meta_attr = sorted(unknown_meta_attrs)[0]
raise TypeError(
add_suggestion(
f"{cls.__name__}() got unexpected keyword '{unknown_meta_attr}'.",
word=unknown_meta_attr,
possibilities=sorted(cls._META_ATTRS),
)
)
if like is not None:
dtype = dtype or like.dtype
device = device or like.device
data = cls._to_tensor(data, dtype=dtype, device=device)
requires_grad = False
self = torch.Tensor._make_subclass(cast(_TensorBase, cls), data, requires_grad)
meta_data = dict()
for attr, (explicit, fallback) in cls._parse_meta_data(**kwargs).items():
if explicit is not DEFAULT:
value = explicit
elif like is not None:
value = getattr(like, attr)
else:
value = fallback(data) if callable(fallback) else fallback
meta_data[attr] = value
self._meta_data = meta_data
return self
for name in meta_attrs:
setattr(cls, name, property(cast(Callable[[F], Any], lambda self, name=name: self._metadata[name])))
def __new__(
cls: Type[F],
data: Any,
*,
dtype: Optional[torch.dtype] = None,
device: Optional[Union[torch.device, str]] = None,
) -> F:
if isinstance(device, str):
device = torch.device(device)
feature = cast(
F,
torch.Tensor._make_subclass(
cast(_TensorBase, cls),
cls._to_tensor(data, dtype=dtype, device=device),
# requires_grad
False,
),
)
feature._metadata = dict()
return feature
@classmethod
def _to_tensor(cls, data, *, dtype, device):
def _to_tensor(self, data: Any, *, dtype: Optional[torch.dtype], device: Optional[torch.device]) -> torch.Tensor:
return torch.as_tensor(data, dtype=dtype, device=device)
@classmethod
def _parse_meta_data(cls) -> Dict[str, Tuple[Any, Any]]:
return dict()
def new_like(
cls: Type[F],
other: F,
data: Any,
*,
dtype: Optional[torch.dtype] = None,
device: Optional[Union[torch.device, str]] = None,
**metadata: Any,
) -> F:
_metadata = other._metadata.copy()
_metadata.update(metadata)
return cls(data, dtype=dtype or other.dtype, device=device or other.device, **_metadata)
@classmethod
def __torch_function__(
......@@ -89,12 +84,37 @@ class Feature(torch.Tensor):
args: Sequence[Any] = (),
kwargs: Optional[Mapping[str, Any]] = None,
) -> torch.Tensor:
"""For general information about how the __torch_function__ protocol works,
see https://pytorch.org/docs/stable/notes/extending.html#extending-torch
TL;DR: Every time a PyTorch operator is called, it goes through the inputs and looks for the
``__torch_function__`` method. If one is found, it is invoked with the operator as ``func`` as well as the
``args`` and ``kwargs`` of the original call.
The default behavior of :class:`~torch.Tensor`'s is to retain a custom tensor type. For the :class:`Feature`
use case, this has two downsides:
1. Since some :class:`Feature`'s require metadata to be constructed, the default wrapping, i.e.
``return cls(func(*args, **kwargs))``, will fail for them.
2. For most operations, there is no way of knowing if the input type is still valid for the output.
For these reasons, the automatic output wrapping is turned off for most operators.
Exceptions to this are:
- :func:`torch.clone`
- :meth:`torch.Tensor.to`
"""
kwargs = kwargs or dict()
with DisableTorchFunction():
output = func(*args, **(kwargs or dict()))
if func is not torch.Tensor.clone:
return output
output = func(*args, **kwargs)
return cls(output, like=args[0])
if func is torch.Tensor.clone:
return cls.new_like(args[0], output)
elif func is torch.Tensor.to:
return cls.new_like(args[0], output, dtype=output.dtype, device=output.device)
else:
return output
def __repr__(self) -> str:
return torch.Tensor.__repr__(self).replace("tensor", type(self).__name__)
return cast(str, torch.Tensor.__repr__(self)).replace("tensor", type(self).__name__)
torchvision/prototype/features/_image.py
View file @ d32bc4ba
from typing import Dict, Any, Union, Tuple
from __future__ import annotations
import warnings
from typing import Any, Optional, Union, Tuple, cast
import torch
from torchvision.prototype.utils._internal import StrEnum
from torchvision.transforms.functional import to_pil_image
from torchvision.utils import draw_bounding_boxes
from torchvision.utils import make_grid
from ._feature import Feature, DEFAULT
from ._bounding_box import BoundingBox
from ._feature import _Feature
class ColorSpace(StrEnum):
# this is just for test purposes
_SENTINEL = -1
OTHER = 0
GRAYSCALE = 1
RGB = 3
class Image(Feature):
color_spaces = ColorSpace
class Image(_Feature):
color_space: ColorSpace
def __new__(
cls,
data: Any,
*,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
color_space: Optional[Union[ColorSpace, str]] = None,
) -> Image:
image = super().__new__(cls, data, dtype=dtype, device=device)
if color_space is None:
color_space = cls.guess_color_space(image)
if color_space == ColorSpace.OTHER:
warnings.warn("Unable to guess a specific color space. Consider passing it explicitly.")
elif isinstance(color_space, str):
color_space = ColorSpace[color_space]
image._metadata.update(dict(color_space=color_space))
return image
@classmethod
def _to_tensor(cls, data, *, dtype, device):
tensor = torch.as_tensor(data, dtype=dtype, device=device)
if tensor.ndim == 2:
def _to_tensor(cls, data: Any, *, dtype: Optional[torch.dtype], device: Optional[torch.device]) -> torch.Tensor:
tensor = super()._to_tensor(data, dtype=dtype, device=device)
if tensor.ndim < 2:
raise ValueError
elif tensor.ndim == 2:
tensor = tensor.unsqueeze(0)
elif tensor.ndim != 3:
raise ValueError("Only single images with 2 or 3 dimensions are allowed.")
return tensor
@classmethod
def _parse_meta_data(
cls,
color_space: Union[str, ColorSpace] = DEFAULT, # type: ignore[assignment]
) -> Dict[str, Tuple[Any, Any]]:
if isinstance(color_space, str):
color_space = ColorSpace[color_space]
return dict(color_space=(color_space, cls.guess_color_space))
@property
def image_size(self) -> Tuple[int, int]:
return cast(Tuple[int, int], self.shape[-2:])
@property
def num_channels(self) -> int:
return self.shape[-3]
@staticmethod
def guess_color_space(data: torch.Tensor) -> ColorSpace:
......@@ -50,3 +74,13 @@ class Image(Feature):
return ColorSpace.RGB
else:
return ColorSpace.OTHER
def show(self) -> None:
# TODO: this is useful for developing and debugging but we should remove or at least revisit this before we
# promote this out of the prototype state
to_pil_image(make_grid(self.view(-1, *self.shape[-3:]))).show()
def draw_bounding_box(self, bounding_box: BoundingBox, **kwargs: Any) -> Image:
# TODO: this is useful for developing and debugging but we should remove or at least revisit this before we
# promote this out of the prototype state
return Image.new_like(self, draw_bounding_boxes(self, bounding_box.to_format("xyxy").view(-1, 4), **kwargs))
torchvision/prototype/features/_label.py
View file @ d32bc4ba
from typing import Dict, Any, Optional, Tuple
from __future__ import annotations
from ._feature import Feature, DEFAULT
from typing import Any, Optional, Sequence, cast
import torch
from torchvision.prototype.utils._internal import apply_recursively
class Label(Feature):
category: Optional[str]
from ._feature import _Feature
class Label(_Feature):
categories: Optional[Sequence[str]]
def __new__(
cls,
data: Any,
*,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
like: Optional[Label] = None,
categories: Optional[Sequence[str]] = None,
) -> Label:
label = super().__new__(cls, data, dtype=dtype, device=device)
label._metadata.update(dict(categories=categories))
return label
@classmethod
def _parse_meta_data(
def from_category(cls, category: str, *, categories: Sequence[str]) -> Label:
return cls(categories.index(category), categories=categories)
def to_categories(self) -> Any:
if not self.categories:
raise RuntimeError()
return apply_recursively(lambda idx: cast(Sequence[str], self.categories)[idx], self.tolist())
class OneHotLabel(_Feature):
categories: Optional[Sequence[str]]
def __new__(
cls,
category: Optional[str] = DEFAULT, # type: ignore[assignment]
) -> Dict[str, Tuple[Any, Any]]:
return dict(category=(category, None))
data: Any,
*,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
like: Optional[Label] = None,
categories: Optional[Sequence[str]] = None,
) -> OneHotLabel:
one_hot_label = super().__new__(cls, data, dtype=dtype, device=device)
if categories is not None and len(categories) != one_hot_label.shape[-1]:
raise ValueError()
one_hot_label._metadata.update(dict(categories=categories))
return one_hot_label
torchvision/prototype/features/_segmentation_mask.py 0 → 100644
View file @ d32bc4ba
from ._feature import _Feature
class SegmentationMask(_Feature):
pass
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