Merge branch 'models/convnext_variants' of https://github.com/datumbox/vision into models/convnext_variants

Author: datumbox

.circleci/config.yml

@@ -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

.circleci/config.yml.in

@@ -23,6 +23,8 @@ supported Python versions.
 +==========================+==========================+=================================+
 | ``main`` / ``nightly``   | ``main`` / ``nightly``   | ``>=3.7``, ``<=3.9``            |
 +--------------------------+--------------------------+---------------------------------+
+| ``1.10.2``               | ``0.11.3``               | ``>=3.6``, ``<=3.9``            |
++--------------------------+--------------------------+---------------------------------+
 | ``1.10.1``               | ``0.11.2``               | ``>=3.6``, ``<=3.9``            |
 +--------------------------+--------------------------+---------------------------------+
 | ``1.10.0``               | ``0.11.1``               | ``>=3.6``, ``<=3.9``            |

README.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

docs/source/utils.rst

@@ -58,6 +58,7 @@ def write_version_file():
     pytorch_dep += "==" + os.getenv("PYTORCH_VERSION")
 
 requirements = [
+    "typing_extensions",
     "numpy",
     "requests",
     pytorch_dep,

setup.py

@@ -2,6 +2,7 @@
 import csv
 import functools
 import gzip
+import io
 import itertools
 import json
 import lzma
@@ -1312,3 +1313,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

test/assets/expected_flow.pt

@@ -18,7 +18,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 ')} "

test/builtin_dataset_mocks.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__])

test/test_prototype_builtin_datasets.py

@@ -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");
   }

test/test_utils.py

@@ -1,5 +1,5 @@
 import os.path
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple, cast
 
 import numpy as np
 from PIL import Image
@@ -65,10 +65,12 @@ def __init__(
         self.labels: Optional[np.ndarray]
         if self.split == "train":
             self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
+            self.labels = cast(np.ndarray, self.labels)
             self.__load_folds(folds)
 
         elif self.split == "train+unlabeled":
             self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
+            self.labels = cast(np.ndarray, self.labels)
             self.__load_folds(folds)
             unlabeled_data, _ = self.__loadfile(self.train_list[2][0])
             self.data = np.concatenate((self.data, unlabeled_data))

torchvision/csrc/io/decoder/gpu/demuxer.h

@@ -6,7 +6,7 @@
 
 from .. import mobilenetv3
 from .. import resnet
-from ..feature_extraction import create_feature_extractor
+from .._utils import IntermediateLayerGetter
 from ._utils import _SimpleSegmentationModel, _load_weights
 from .fcn import FCNHead
 
@@ -121,7 +121,7 @@ def _deeplabv3_resnet(
     return_layers = {"layer4": "out"}
     if aux:
         return_layers["layer3"] = "aux"
-    backbone = create_feature_extractor(backbone, return_layers)
+    backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
 
     aux_classifier = FCNHead(1024, num_classes) if aux else None
     classifier = DeepLabHead(2048, num_classes)
@@ -144,7 +144,7 @@ def _deeplabv3_mobilenetv3(
     return_layers = {str(out_pos): "out"}
     if aux:
         return_layers[str(aux_pos)] = "aux"
-    backbone = create_feature_extractor(backbone, return_layers)
+    backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
 
     aux_classifier = FCNHead(aux_inplanes, num_classes) if aux else None
     classifier = DeepLabHead(out_inplanes, num_classes)

torchvision/datasets/stl10.py

@@ -3,7 +3,7 @@
 from torch import nn
 
 from .. import resnet
-from ..feature_extraction import create_feature_extractor
+from .._utils import IntermediateLayerGetter
 from ._utils import _SimpleSegmentationModel, _load_weights
 
 
@@ -57,7 +57,7 @@ def _fcn_resnet(
     return_layers = {"layer4": "out"}
     if aux:
         return_layers["layer3"] = "aux"
-    backbone = create_feature_extractor(backbone, return_layers)
+    backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)
 
     aux_classifier = FCNHead(1024, num_classes) if aux else None
     classifier = FCNHead(2048, num_classes)

torchvision/models/segmentation/deeplabv3.py

@@ -6,7 +6,7 @@
 
 from ...utils import _log_api_usage_once
 from .. import mobilenetv3
-from ..feature_extraction import create_feature_extractor
+from .._utils import IntermediateLayerGetter
 from ._utils import _load_weights
 
 
@@ -90,7 +90,7 @@ def _lraspp_mobilenetv3(backbone: mobilenetv3.MobileNetV3, num_classes: int) ->
     high_pos = stage_indices[-1]  # use C5 which has output_stride = 16
     low_channels = backbone[low_pos].out_channels
     high_channels = backbone[high_pos].out_channels
-    backbone = create_feature_extractor(backbone, {str(low_pos): "low", str(high_pos): "high"})
+    backbone = IntermediateLayerGetter(backbone, return_layers={str(low_pos): "low", str(high_pos): "high"})
 
     return LRASPP(backbone, low_channels, high_channels, num_classes)
 

torchvision/models/segmentation/fcn.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

torchvision/models/segmentation/lraspp.py

@@ -23,8 +23,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())
 
 
@@ -39,7 +38,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

torchvision/prototype/datasets/__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

torchvision/prototype/datasets/_api.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)

torchvision/prototype/datasets/_builtin/__init__.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
 

torchvision/prototype/datasets/_builtin/pcam.py

@@ -24,6 +24,8 @@ class DatasetType(enum.Enum):
 
 
 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