Merge branch 'main' into revamp-prototype-features-transforms

pmeier · pmeier · commit 3e23333e1840 · 2022-01-27T15:29:30.000+01:00
diff --git a/.circleci/config.yml b/.circleci/config.yml
diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in
@@ -351,7 +351,7 @@ jobs:
       - install_torchvision
       - install_prototype_dependencies
       - pip_install:
-          args: scipy pycocotools
+          args: scipy pycocotools h5py
           descr: Install optional dependencies
       - run:
           name: Enable prototype tests
diff --git a/docs/source/utils.rst b/docs/source/utils.rst
@@ -15,5 +15,6 @@ vizualization <sphx_glr_auto_examples_plot_visualization_utils.py>`.
     draw_bounding_boxes
     draw_segmentation_masks
     draw_keypoints
+    flow_to_image
     make_grid
     save_image
diff --git a/references/detection/coco_utils.py b/references/detection/coco_utils.py
@@ -156,7 +156,7 @@ def convert_to_coco_api(ds):
         img_dict["height"] = img.shape[-2]
         img_dict["width"] = img.shape[-1]
         dataset["images"].append(img_dict)
-        bboxes = targets["boxes"]
+        bboxes = targets["boxes"].clone()
         bboxes[:, 2:] -= bboxes[:, :2]
         bboxes = bboxes.tolist()
         labels = targets["labels"].tolist()
diff --git a/test/assets/expected_flow.pt b/test/assets/expected_flow.pt
diff --git a/test/builtin_dataset_mocks.py b/test/builtin_dataset_mocks.py
@@ -2,6 +2,7 @@
 import csv
 import functools
 import gzip
+import io
 import itertools
 import json
 import lzma
@@ -1316,3 +1317,30 @@ def svhn(info, root, config):
         },
     )
     return num_samples
+
+
+@register_mock
+def pcam(info, root, config):
+    import h5py
+
+    num_images = {"train": 2, "test": 3, "val": 4}[config.split]
+
+    split = "valid" if config.split == "val" else config.split
+
+    images_io = io.BytesIO()
+    with h5py.File(images_io, "w") as f:
+        f["x"] = np.random.randint(0, 256, size=(num_images, 10, 10, 3), dtype=np.uint8)
+
+    targets_io = io.BytesIO()
+    with h5py.File(targets_io, "w") as f:
+        f["y"] = np.random.randint(0, 2, size=(num_images, 1, 1, 1), dtype=np.uint8)
+
+    # Create .gz compressed files
+    images_file = root / f"camelyonpatch_level_2_split_{split}_x.h5.gz"
+    targets_file = root / f"camelyonpatch_level_2_split_{split}_y.h5.gz"
+    for compressed_file_name, uncompressed_file_io in ((images_file, images_io), (targets_file, targets_io)):
+        compressed_data = gzip.compress(uncompressed_file_io.getbuffer())
+        with open(compressed_file_name, "wb") as compressed_file:
+            compressed_file.write(compressed_data)
+
+    return num_images
diff --git a/test/test_prototype_builtin_datasets.py b/test/test_prototype_builtin_datasets.py
@@ -25,7 +25,7 @@ def test_home(mocker, tmp_path):
 
 
 def test_coverage():
-    untested_datasets = set(datasets.list()) - DATASET_MOCKS.keys()
+    untested_datasets = set(datasets.list_datasets()) - DATASET_MOCKS.keys()
     if untested_datasets:
         raise AssertionError(
             f"The dataset(s) {sequence_to_str(sorted(untested_datasets), separate_last='and ')} "
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -317,5 +317,30 @@ def test_draw_keypoints_errors():
         utils.draw_keypoints(image=img, keypoints=invalid_keypoints)
 
 
+def test_flow_to_image():
+    h, w = 100, 100
+    flow = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    flow = torch.stack(flow[::-1], dim=0).float()
+    flow[0] -= h / 2
+    flow[1] -= w / 2
+    img = utils.flow_to_image(flow)
+    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "expected_flow.pt")
+    expected_img = torch.load(path, map_location="cpu")
+    assert_equal(expected_img, img)
+
+
+def test_flow_to_image_errors():
+    wrong_flow1 = torch.full((3, 10, 10), 0, dtype=torch.float)
+    wrong_flow2 = torch.full((2, 10), 0, dtype=torch.float)
+    wrong_flow3 = torch.full((2, 10, 30), 0, dtype=torch.int)
+
+    with pytest.raises(ValueError, match="Input flow should have shape"):
+        utils.flow_to_image(flow=wrong_flow1)
+    with pytest.raises(ValueError, match="Input flow should have shape"):
+        utils.flow_to_image(flow=wrong_flow2)
+    with pytest.raises(ValueError, match="Flow should be of dtype torch.float"):
+        utils.flow_to_image(flow=wrong_flow3)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/torchvision/csrc/io/decoder/gpu/demuxer.h b/torchvision/csrc/io/decoder/gpu/demuxer.h
@@ -119,6 +119,7 @@ class Demuxer {
           " in demuxer.h\n");
     }
   }
+
   ~Demuxer() {
     if (!fmtCtx) {
       return;
@@ -223,7 +224,7 @@ class Demuxer {
     int64_t time = timestamp * AV_TIME_BASE;
     TORCH_CHECK(
         0 <= av_seek_frame(fmtCtx, -1, time, flag),
-        "avformat_open_input() failed at line ",
+        "av_seek_frame() failed at line ",
         __LINE__,
         " in demuxer.h\n");
   }
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 as list, info, load, find  # usort: skip
+from ._api import register, list_datasets, info, load, find  # 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
@@ -20,8 +20,7 @@ def register(dataset: Dataset) -> None:
         register(obj())
 
 
-# This is exposed as 'list', but we avoid that here to not shadow the built-in 'list'
-def _list() -> List[str]:
+def list_datasets() -> List[str]:
     return sorted(DATASETS.keys())
 
 
@@ -36,7 +35,7 @@ def find(name: str) -> Dataset:
                 word=name,
                 possibilities=DATASETS.keys(),
                 alternative_hint=lambda _: (
-                    "You can use torchvision.datasets.list() to get a list of all available datasets."
+                    "You can use torchvision.datasets.list_datasets() to get a list of all available datasets."
                 ),
             )
         ) from error
diff --git a/torchvision/prototype/datasets/_builtin/__init__.py b/torchvision/prototype/datasets/_builtin/__init__.py
@@ -10,6 +10,7 @@
 from .imagenet import ImageNet
 from .mnist import MNIST, FashionMNIST, KMNIST, EMNIST, QMNIST
 from .oxford_iiit_pet import OxfordIITPet
+from .pcam import PCAM
 from .sbd import SBD
 from .semeion import SEMEION
 from .svhn import SVHN
diff --git a/torchvision/prototype/datasets/_builtin/pcam.py b/torchvision/prototype/datasets/_builtin/pcam.py
@@ -0,0 +1,125 @@
+import io
+from collections import namedtuple
+from typing import Any, Callable, 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 (
+    Dataset,
+    DatasetConfig,
+    DatasetInfo,
+    OnlineResource,
+    DatasetType,
+    GDriveResource,
+)
+from torchvision.prototype.datasets.utils._internal import (
+    hint_sharding,
+    hint_shuffling,
+)
+from torchvision.prototype.features import Label
+
+
+class PCAMH5Reader(IterDataPipe[Tuple[str, io.IOBase]]):
+    def __init__(
+        self,
+        datapipe: IterDataPipe[Tuple[str, io.IOBase]],
+        key: Optional[str] = None,  # Note: this key thing might be very specific to the PCAM dataset
+    ) -> None:
+        self.datapipe = datapipe
+        self.key = key
+
+    def __iter__(self) -> Iterator[Tuple[str, io.IOBase]]:
+        import h5py
+
+        for _, handle in self.datapipe:
+            with h5py.File(handle) as data:
+                if self.key is not None:
+                    data = data[self.key]
+                yield from data
+
+
+_Resource = namedtuple("_Resource", ("file_name", "gdrive_id", "sha256"))
+
+
+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")),
+            dependencies=["h5py"],
+        )
+
+    _RESOURCES = {
+        "train": (
+            _Resource(  # Images
+                file_name="camelyonpatch_level_2_split_train_x.h5.gz",
+                gdrive_id="1Ka0XfEMiwgCYPdTI-vv6eUElOBnKFKQ2",
+                sha256="d619e741468a7ab35c7e4a75e6821b7e7e6c9411705d45708f2a0efc8960656c",
+            ),
+            _Resource(  # Targets
+                file_name="camelyonpatch_level_2_split_train_y.h5.gz",
+                gdrive_id="1269yhu3pZDP8UYFQs-NYs3FPwuK-nGSG",
+                sha256="b74126d2c01b20d3661f9b46765d29cf4e4fba6faba29c8e0d09d406331ab75a",
+            ),
+        ),
+        "test": (
+            _Resource(  # Images
+                file_name="camelyonpatch_level_2_split_test_x.h5.gz",
+                gdrive_id="1qV65ZqZvWzuIVthK8eVDhIwrbnsJdbg_",
+                sha256="79174c2201ad521602a5888be8f36ee10875f37403dd3f2086caf2182ef87245",
+            ),
+            _Resource(  # Targets
+                file_name="camelyonpatch_level_2_split_test_y.h5.gz",
+                gdrive_id="17BHrSrwWKjYsOgTMmoqrIjDy6Fa2o_gP",
+                sha256="0a522005fccc8bbd04c5a117bfaf81d8da2676f03a29d7499f71d0a0bd6068ef",
+            ),
+        ),
+        "val": (
+            _Resource(  # Images
+                file_name="camelyonpatch_level_2_split_valid_x.h5.gz",
+                gdrive_id="1hgshYGWK8V-eGRy8LToWJJgDU_rXWVJ3",
+                sha256="f82ee1670d027b4ec388048d9eabc2186b77c009655dae76d624c0ecb053ccb2",
+            ),
+            _Resource(  # Targets
+                file_name="camelyonpatch_level_2_split_valid_y.h5.gz",
+                gdrive_id="1bH8ZRbhSVAhScTS0p9-ZzGnX91cHT3uO",
+                sha256="ce1ae30f08feb468447971cfd0472e7becd0ad96d877c64120c72571439ae48c",
+            ),
+        ),
+    }
+
+    def resources(self, config: DatasetConfig) -> List[OnlineResource]:
+        return [  # = [images resource, targets resource]
+            GDriveResource(file_name=file_name, id=gdrive_id, sha256=sha256, decompress=True)
+            for file_name, gdrive_id, sha256 in self._RESOURCES[config.split]
+        ]
+
+    def _collate_and_decode(self, data: Tuple[Any, Any]) -> Dict[str, Any]:
+        image, target = data  # They're both numpy arrays at this point
+
+        return {
+            "image": features.Image(image),
+            "label": Label(target.item()),
+        }
+
+    def _make_datapipe(
+        self,
+        resource_dps: List[IterDataPipe],
+        *,
+        config: DatasetConfig,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> IterDataPipe[Dict[str, Any]]:
+
+        images_dp, targets_dp = resource_dps
+
+        images_dp = PCAMH5Reader(images_dp, key="x")
+        targets_dp = PCAMH5Reader(targets_dp, key="y")
+
+        dp = Zipper(images_dp, targets_dp)
+        dp = hint_sharding(dp)
+        dp = hint_shuffling(dp)
+        return Mapper(dp, self._collate_and_decode)
diff --git a/torchvision/prototype/datasets/generate_category_files.py b/torchvision/prototype/datasets/generate_category_files.py
@@ -49,7 +49,7 @@ def parse_args(argv=None):
     args = parser.parse_args(argv or sys.argv[1:])
 
     if not args.names:
-        args.names = datasets.list()
+        args.names = datasets.list_datasets()
 
     return args
 
diff --git a/torchvision/prototype/datasets/utils/_dataset.py b/torchvision/prototype/datasets/utils/_dataset.py
@@ -15,6 +15,8 @@
 
 
 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.
     pass
 
 
diff --git a/torchvision/prototype/datasets/utils/_internal.py b/torchvision/prototype/datasets/utils/_internal.py
@@ -34,7 +34,6 @@
     "INFINITE_BUFFER_SIZE",
     "BUILTIN_DIR",
     "read_mat",
-    "SequenceIterator",
     "MappingIterator",
     "Enumerator",
     "getitem",
@@ -66,15 +65,6 @@ def read_mat(buffer: BinaryIO, **kwargs: Any) -> Any:
     return sio.loadmat(buffer, **kwargs)
 
 
-class SequenceIterator(IterDataPipe[D]):
-    def __init__(self, datapipe: IterDataPipe[Sequence[D]]):
-        self.datapipe = datapipe
-
-    def __iter__(self) -> Iterator[D]:
-        for sequence in self.datapipe:
-            yield from iter(sequence)
-
-
 class MappingIterator(IterDataPipe[Union[Tuple[K, D], D]]):
     def __init__(self, datapipe: IterDataPipe[Dict[K, D]], *, drop_key: bool = False) -> None:
         self.datapipe = datapipe
diff --git a/torchvision/prototype/datasets/utils/_resource.py b/torchvision/prototype/datasets/utils/_resource.py
@@ -2,7 +2,6 @@
 import hashlib
 import itertools
 import pathlib
-import warnings
 from typing import Optional, Sequence, Tuple, Callable, IO, Any, Union, NoReturn
 from urllib.parse import urlparse
 
@@ -32,23 +31,17 @@ def __init__(
         sha256: Optional[str] = None,
         decompress: bool = False,
         extract: bool = False,
-        preprocess: Optional[Callable[[pathlib.Path], pathlib.Path]] = None,
-        loader: Optional[Callable[[pathlib.Path], IterDataPipe[Tuple[str, IO]]]] = None,
     ) -> None:
         self.file_name = file_name
         self.sha256 = sha256
 
-        if preprocess and (decompress or extract):
-            warnings.warn("The parameters 'decompress' and 'extract' are ignored when 'preprocess' is passed.")
-        elif extract:
-            preprocess = self._extract
+        self._preprocess: Optional[Callable[[pathlib.Path], pathlib.Path]]
+        if extract:
+            self._preprocess = self._extract
         elif decompress:
-            preprocess = self._decompress
-        self._preprocess = preprocess
-
-        if loader is None:
-            loader = self._default_loader
-        self._loader = loader
+            self._preprocess = self._decompress
+        else:
+            self._preprocess = None
 
     @staticmethod
     def _extract(file: pathlib.Path) -> pathlib.Path:
@@ -60,7 +53,7 @@ def _extract(file: pathlib.Path) -> pathlib.Path:
     def _decompress(file: pathlib.Path) -> pathlib.Path:
         return pathlib.Path(_decompress(str(file), remove_finished=True))
 
-    def _default_loader(self, path: pathlib.Path) -> IterDataPipe[Tuple[str, IO]]:
+    def _loader(self, path: pathlib.Path) -> IterDataPipe[Tuple[str, IO]]:
         if path.is_dir():
             return FileOpener(FileLister(str(path), recursive=True), mode="rb")
 
@@ -101,7 +94,7 @@ def load(
             path_candidates = {self.download(root, skip_integrity_check=skip_integrity_check)}
         # If the only thing we find is the raw file, we use it and optionally perform some preprocessing steps.
         if path_candidates == {path}:
-            if self._preprocess:
+            if self._preprocess is not None:
                 path = self._preprocess(path)
         # Otherwise we use the path with the fewest suffixes. This gives us the extracted > decompressed > raw priority
         # that we want.
diff --git a/torchvision/prototype/models/quantization/mobilenetv3.py b/torchvision/prototype/models/quantization/mobilenetv3.py
diff --git a/torchvision/utils.py b/torchvision/utils.py

Original file line number	Diff line number	Diff line change
`@@ -119,6 +119,7 @@ class Demuxer {`
`119`	`119`	`" in demuxer.h\n");`
`120`	`120`	`}`
`121`	`121`	`}`
	`122`	`+`
`122`	`123`	`~Demuxer() {`
`123`	`124`	`if (!fmtCtx) {`
`124`	`125`	`return;`
`@@ -223,7 +224,7 @@ class Demuxer {`
`223`	`224`	`int64_t time = timestamp * AV_TIME_BASE;`
`224`	`225`	`TORCH_CHECK(`
`225`	`226`	`0 <= av_seek_frame(fmtCtx, -1, time, flag),`
`226`		`- "avformat_open_input() failed at line ",`
	`227`	`+ "av_seek_frame() failed at line ",`
`227`	`228`	`__LINE__,`
`228`	`229`	`" in demuxer.h\n");`
`229`	`230`	`}`