Migrate Food101 prototype dataset (#5758)

* Migrate Food101 dataset

* Added length

* Update torchvision/prototype/datasets/_builtin/food101.py

Co-authored-by: Philip Meier <[email protected]>

Co-authored-by: Philip Meier <[email protected]>

Author: NicolasHug

test/builtin_dataset_mocks.py

@@ -933,8 +933,8 @@ def country211(root, config):
     return num_examples * len(classes)
 
 
-# @register_mock
-def food101(info, root, config):
+@register_mock(configs=combinations_grid(split=("train", "test")))
+def food101(root, config):
     data_folder = root / "food-101"
 
     num_images_per_class = 3
@@ -968,7 +968,7 @@ def food101(info, root, config):
 
     make_tar(root, f"{data_folder.name}.tar.gz", compression="gz")
 
-    return num_samples_map[config.split]
+    return num_samples_map[config["split"]]
 
 
 @register_mock(configs=combinations_grid(split=("train", "val", "test"), fold=(1, 4, 10)))

torchvision/prototype/datasets/_builtin/food101.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Tuple, List, Dict, Optional, BinaryIO
+from typing import Any, Tuple, List, Dict, Optional, BinaryIO, Union
 
 from torchdata.datapipes.iter import (
     IterDataPipe,
@@ -9,26 +9,43 @@
     Demultiplexer,
     IterKeyZipper,
 )
-from torchvision.prototype.datasets.utils import Dataset, DatasetInfo, DatasetConfig, HttpResource, OnlineResource
+from torchvision.prototype.datasets.utils import Dataset2, DatasetInfo, HttpResource, OnlineResource
 from torchvision.prototype.datasets.utils._internal import (
     hint_shuffling,
+    BUILTIN_DIR,
     hint_sharding,
     path_comparator,
     getitem,
     INFINITE_BUFFER_SIZE,
 )
 from torchvision.prototype.features import Label, EncodedImage
 
+from .._api import register_dataset, register_info
+
+
+NAME = "food101"
+
+
+@register_info(NAME)
+def _info() -> Dict[str, Any]:
+    categories = DatasetInfo.read_categories_file(BUILTIN_DIR / f"{NAME}.categories")
+    categories = [c[0] for c in categories]
+    return dict(categories=categories)
 
-class Food101(Dataset):
-    def _make_info(self) -> DatasetInfo:
-        return DatasetInfo(
-            "food101",
-            homepage="https://data.vision.ee.ethz.ch/cvl/datasets_extra/food-101",
-            valid_options=dict(split=("train", "test")),
-        )
 
-    def resources(self, config: DatasetConfig) -> List[OnlineResource]:
+@register_dataset(NAME)
+class Food101(Dataset2):
+    """Food 101 dataset
+    homepage="https://data.vision.ee.ethz.ch/cvl/datasets_extra/food-101",
+    """
+
+    def __init__(self, root: Union[str, Path], *, split: str = "train", skip_integrity_check: bool = False) -> None:
+        self._split = self._verify_str_arg(split, "split", {"train", "test"})
+        self._categories = _info()["categories"]
+
+        super().__init__(root, skip_integrity_check=skip_integrity_check)
+
+    def _resources(self) -> List[OnlineResource]:
         return [
             HttpResource(
                 url="http://data.vision.ee.ethz.ch/cvl/food-101.tar.gz",
@@ -49,7 +66,7 @@ def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
     def _prepare_sample(self, data: Tuple[str, Tuple[str, BinaryIO]]) -> Dict[str, Any]:
         id, (path, buffer) = data
         return dict(
-            label=Label.from_category(id.split("/", 1)[0], categories=self.categories),
+            label=Label.from_category(id.split("/", 1)[0], categories=self._categories),
             path=path,
             image=EncodedImage.from_file(buffer),
         )
@@ -58,17 +75,12 @@ def _image_key(self, data: Tuple[str, Any]) -> str:
         path = Path(data[0])
         return path.relative_to(path.parents[1]).with_suffix("").as_posix()
 
-    def _make_datapipe(
-        self,
-        resource_dps: List[IterDataPipe],
-        *,
-        config: DatasetConfig,
-    ) -> IterDataPipe[Dict[str, Any]]:
+    def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:
         archive_dp = resource_dps[0]
         images_dp, split_dp = Demultiplexer(
             archive_dp, 2, self._classify_archive, drop_none=True, buffer_size=INFINITE_BUFFER_SIZE
         )
-        split_dp = Filter(split_dp, path_comparator("name", f"{config.split}.txt"))
+        split_dp = Filter(split_dp, path_comparator("name", f"{self._split}.txt"))
         split_dp = LineReader(split_dp, decode=True, return_path=False)
         split_dp = hint_sharding(split_dp)
         split_dp = hint_shuffling(split_dp)
@@ -83,9 +95,12 @@ def _make_datapipe(
 
         return Mapper(dp, self._prepare_sample)
 
-    def _generate_categories(self, root: Path) -> List[str]:
-        resources = self.resources(self.default_config)
-        dp = resources[0].load(root)
+    def _generate_categories(self) -> List[str]:
+        resources = self.resources()
+        dp = resources[0].load(self._root)
         dp = Filter(dp, path_comparator("name", "classes.txt"))
         dp = LineReader(dp, decode=True, return_path=False)
         return list(dp)
+
+    def __len__(self) -> int:
+        return 75_750 if self._split == "train" else 25_250