add CUB200 prototype datasets (#5154)

* add CUB200 prototype datasets

* address review comments

torchvision/prototype/datasets/_builtin/__init__.py

@@ -3,6 +3,7 @@ from .celeba import CelebA
 from .cifar import Cifar10, Cifar100
 from .clevr import CLEVR
 from .coco import Coco
+from .cub200 import CUB200
 from .dtd import DTD
 from .fer2013 import FER2013
 from .imagenet import ImageNet

torchvision/prototype/datasets/_builtin/cub200.categories

+Black_footed_Albatross
+Laysan_Albatross
+Sooty_Albatross
+Groove_billed_Ani
+Crested_Auklet
+Least_Auklet
+Parakeet_Auklet
+Rhinoceros_Auklet
+Brewer_Blackbird
+Red_winged_Blackbird
+Rusty_Blackbird
+Yellow_headed_Blackbird
+Bobolink
+Indigo_Bunting
+Lazuli_Bunting
+Painted_Bunting
+Cardinal
+Spotted_Catbird
+Gray_Catbird
+Yellow_breasted_Chat
+Eastern_Towhee
+Chuck_will_Widow
+Brandt_Cormorant
+Red_faced_Cormorant
+Pelagic_Cormorant
+Bronzed_Cowbird
+Shiny_Cowbird
+Brown_Creeper
+American_Crow
+Fish_Crow
+Black_billed_Cuckoo
+Mangrove_Cuckoo
+Yellow_billed_Cuckoo
+Gray_crowned_Rosy_Finch
+Purple_Finch
+Northern_Flicker
+Acadian_Flycatcher
+Great_Crested_Flycatcher
+Least_Flycatcher
+Olive_sided_Flycatcher
+Scissor_tailed_Flycatcher
+Vermilion_Flycatcher
+Yellow_bellied_Flycatcher
+Frigatebird
+Northern_Fulmar
+Gadwall
+American_Goldfinch
+European_Goldfinch
+Boat_tailed_Grackle
+Eared_Grebe
+Horned_Grebe
+Pied_billed_Grebe
+Western_Grebe
+Blue_Grosbeak
+Evening_Grosbeak
+Pine_Grosbeak
+Rose_breasted_Grosbeak
+Pigeon_Guillemot
+California_Gull
+Glaucous_winged_Gull
+Heermann_Gull
+Herring_Gull
+Ivory_Gull
+Ring_billed_Gull
+Slaty_backed_Gull
+Western_Gull
+Anna_Hummingbird
+Ruby_throated_Hummingbird
+Rufous_Hummingbird
+Green_Violetear
+Long_tailed_Jaeger
+Pomarine_Jaeger
+Blue_Jay
+Florida_Jay
+Green_Jay
+Dark_eyed_Junco
+Tropical_Kingbird
+Gray_Kingbird
+Belted_Kingfisher
+Green_Kingfisher
+Pied_Kingfisher
+Ringed_Kingfisher
+White_breasted_Kingfisher
+Red_legged_Kittiwake
+Horned_Lark
+Pacific_Loon
+Mallard
+Western_Meadowlark
+Hooded_Merganser
+Red_breasted_Merganser
+Mockingbird
+Nighthawk
+Clark_Nutcracker
+White_breasted_Nuthatch
+Baltimore_Oriole
+Hooded_Oriole
+Orchard_Oriole
+Scott_Oriole
+Ovenbird
+Brown_Pelican
+White_Pelican
+Western_Wood_Pewee
+Sayornis
+American_Pipit
+Whip_poor_Will
+Horned_Puffin
+Common_Raven
+White_necked_Raven
+American_Redstart
+Geococcyx
+Loggerhead_Shrike
+Great_Grey_Shrike
+Baird_Sparrow
+Black_throated_Sparrow
+Brewer_Sparrow
+Chipping_Sparrow
+Clay_colored_Sparrow
+House_Sparrow
+Field_Sparrow
+Fox_Sparrow
+Grasshopper_Sparrow
+Harris_Sparrow
+Henslow_Sparrow
+Le_Conte_Sparrow
+Lincoln_Sparrow
+Nelson_Sharp_tailed_Sparrow
+Savannah_Sparrow
+Seaside_Sparrow
+Song_Sparrow
+Tree_Sparrow
+Vesper_Sparrow
+White_crowned_Sparrow
+White_throated_Sparrow
+Cape_Glossy_Starling
+Bank_Swallow
+Barn_Swallow
+Cliff_Swallow
+Tree_Swallow
+Scarlet_Tanager
+Summer_Tanager
+Artic_Tern
+Black_Tern
+Caspian_Tern
+Common_Tern
+Elegant_Tern
+Forsters_Tern
+Least_Tern
+Green_tailed_Towhee
+Brown_Thrasher
+Sage_Thrasher
+Black_capped_Vireo
+Blue_headed_Vireo
+Philadelphia_Vireo
+Red_eyed_Vireo
+Warbling_Vireo
+White_eyed_Vireo
+Yellow_throated_Vireo
+Bay_breasted_Warbler
+Black_and_white_Warbler
+Black_throated_Blue_Warbler
+Blue_winged_Warbler
+Canada_Warbler
+Cape_May_Warbler
+Cerulean_Warbler
+Chestnut_sided_Warbler
+Golden_winged_Warbler
+Hooded_Warbler
+Kentucky_Warbler
+Magnolia_Warbler
+Mourning_Warbler
+Myrtle_Warbler
+Nashville_Warbler
+Orange_crowned_Warbler
+Palm_Warbler
+Pine_Warbler
+Prairie_Warbler
+Prothonotary_Warbler
+Swainson_Warbler
+Tennessee_Warbler
+Wilson_Warbler
+Worm_eating_Warbler
+Yellow_Warbler
+Northern_Waterthrush
+Louisiana_Waterthrush
+Bohemian_Waxwing
+Cedar_Waxwing
+American_Three_toed_Woodpecker
+Pileated_Woodpecker
+Red_bellied_Woodpecker
+Red_cockaded_Woodpecker
+Red_headed_Woodpecker
+Downy_Woodpecker
+Bewick_Wren
+Cactus_Wren
+Carolina_Wren
+House_Wren
+Marsh_Wren
+Rock_Wren
+Winter_Wren
+Common_Yellowthroat

torchvision/prototype/datasets/_builtin/cub200.py

+import csv
+import functools
+import io
+import pathlib
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+from torchdata.datapipes.iter import (
+    IterDataPipe,
+    Mapper,
+    Filter,
+    IterKeyZipper,
+    Demultiplexer,
+    LineReader,
+    CSVParser,
+    CSVDictParser,
+)
+from torchvision.prototype.datasets.utils import (
+    Dataset,
+    DatasetConfig,
+    DatasetInfo,
+    HttpResource,
+    OnlineResource,
+    DatasetType,
+)
+from torchvision.prototype.datasets.utils._internal import (
+    INFINITE_BUFFER_SIZE,
+    read_mat,
+    hint_sharding,
+    hint_shuffling,
+    getitem,
+    path_comparator,
+    path_accessor,
+)
+from torchvision.prototype.features import Label, BoundingBox, Feature
+
+csv.register_dialect("cub200", delimiter=" ")
+
+
+class CUB200(Dataset):
+    def _make_info(self) -> DatasetInfo:
+        return DatasetInfo(
+            "cub200",
+            type=DatasetType.IMAGE,
+            homepage="http://www.vision.caltech.edu/visipedia/CUB-200-2011.html",
+            dependencies=("scipy",),
+            valid_options=dict(
+                split=("train", "test"),
+                year=("2011", "2010"),
+            ),
+        )
+
+    def resources(self, config: DatasetConfig) -> List[OnlineResource]:
+        if config.year == "2011":
+            archive = HttpResource(
+                "http://www.vision.caltech.edu/visipedia-data/CUB-200-2011/CUB_200_2011.tgz",
+                sha256="0c685df5597a8b24909f6a7c9db6d11e008733779a671760afef78feb49bf081",
+                decompress=True,
+            )
+            segmentations = HttpResource(
+                "http://www.vision.caltech.edu/visipedia-data/CUB-200-2011/segmentations.tgz",
+                sha256="dc77f6cffea0cbe2e41d4201115c8f29a6320ecb04fffd2444f51b8066e4b84f",
+                decompress=True,
+            )
+            return [archive, segmentations]
+        else:  # config.year == "2010"
+            split = HttpResource(
+                "http://www.vision.caltech.edu/visipedia-data/CUB-200/lists.tgz",
+                sha256="aeacbd5e3539ae84ea726e8a266a9a119c18f055cd80f3836d5eb4500b005428",
+                decompress=True,
+            )
+            images = HttpResource(
+                "http://www.vision.caltech.edu/visipedia-data/CUB-200/images.tgz",
+                sha256="2a6d2246bbb9778ca03aa94e2e683ccb4f8821a36b7f235c0822e659d60a803e",
+                decompress=True,
+            )
+            anns = HttpResource(
+                "http://www.vision.caltech.edu/visipedia-data/CUB-200/annotations.tgz",
+                sha256="c17b7841c21a66aa44ba8fe92369cc95dfc998946081828b1d7b8a4b716805c1",
+                decompress=True,
+            )
+            return [split, images, anns]
+
+    def _2011_classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
+        path = pathlib.Path(data[0])
+        if path.parents[1].name == "images":
+            return 0
+        elif path.name == "train_test_split.txt":
+            return 1
+        elif path.name == "images.txt":
+            return 2
+        elif path.name == "bounding_boxes.txt":
+            return 3
+        else:
+            return None
+
+    def _2011_filter_split(self, row: List[str], *, split: str) -> bool:
+        _, split_id = row
+        return {
+            "0": "test",
+            "1": "train",
+        }[split_id] == split
+
+    def _2011_segmentation_key(self, data: Tuple[str, Any]) -> str:
+        path = pathlib.Path(data[0])
+        return path.with_suffix(".jpg").name
+
+    def _2011_decode_ann(
+        self,
+        data: Tuple[str, Tuple[List[str], Tuple[str, io.IOBase]]],
+        *,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> Dict[str, Any]:
+        _, (bounding_box_data, segmentation_data) = data
+        segmentation_path, segmentation_buffer = segmentation_data
+        return dict(
+            bounding_box=BoundingBox([float(part) for part in bounding_box_data[1:]], format="xywh"),
+            segmentation_path=segmentation_path,
+            segmentation=Feature(decoder(segmentation_buffer)) if decoder else segmentation_buffer,
+        )
+
+    def _2010_split_key(self, data: str) -> str:
+        return data.rsplit("/", maxsplit=1)[1]
+
+    def _2010_anns_key(self, data: Tuple[str, io.IOBase]) -> Tuple[str, Tuple[str, io.IOBase]]:
+        path = pathlib.Path(data[0])
+        return path.with_suffix(".jpg").name, data
+
+    def _2010_decode_ann(
+        self, data: Tuple[str, Tuple[str, io.IOBase]], *, decoder: Optional[Callable[[io.IOBase], torch.Tensor]]
+    ) -> Dict[str, Any]:
+        _, (path, buffer) = data
+        content = read_mat(buffer)
+        return dict(
+            ann_path=path,
+            bounding_box=BoundingBox(
+                [int(content["bbox"][coord]) for coord in ("left", "bottom", "right", "top")], format="xyxy"
+            ),
+            segmentation=Feature(content["seg"]),
+        )
+
+    def _collate_and_decode_sample(
+        self,
+        data: Tuple[Tuple[str, Tuple[str, io.IOBase]], Any],
+        *,
+        year: str,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> Dict[str, Any]:
+        data, anns_data = data
+        _, image_data = data
+        path, buffer = image_data
+
+        dir_name = pathlib.Path(path).parent.name
+        label_str, category = dir_name.split(".")
+
+        return dict(
+            (self._2011_decode_ann if year == "2011" else self._2010_decode_ann)(anns_data, decoder=decoder),
+            image=decoder(buffer) if decoder else buffer,
+            label=Label(int(label_str), category=category),
+        )
+
+    def _make_datapipe(
+        self,
+        resource_dps: List[IterDataPipe],
+        *,
+        config: DatasetConfig,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> IterDataPipe[Dict[str, Any]]:
+        if config.year == "2011":
+            archive_dp, segmentations_dp = resource_dps
+            images_dp, split_dp, image_files_dp, bounding_boxes_dp = Demultiplexer(
+                archive_dp, 4, self._2011_classify_archive, drop_none=True, buffer_size=INFINITE_BUFFER_SIZE
+            )
+
+            image_files_dp = CSVParser(image_files_dp, dialect="cub200")
+            image_files_map = dict(
+                (image_id, rel_posix_path.rsplit("/", maxsplit=1)[1]) for image_id, rel_posix_path in image_files_dp
+            )
+
+            split_dp = CSVParser(split_dp, dialect="cub200")
+            split_dp = Filter(split_dp, functools.partial(self._2011_filter_split, split=config.split))
+            split_dp = Mapper(split_dp, getitem(0))
+            split_dp = Mapper(split_dp, image_files_map.get)
+
+            bounding_boxes_dp = CSVParser(bounding_boxes_dp, dialect="cub200")
+            bounding_boxes_dp = Mapper(bounding_boxes_dp, image_files_map.get, input_col=0)
+
+            anns_dp = IterKeyZipper(
+                bounding_boxes_dp,
+                segmentations_dp,
+                key_fn=getitem(0),
+                ref_key_fn=self._2011_segmentation_key,
+                keep_key=True,
+                buffer_size=INFINITE_BUFFER_SIZE,
+            )
+        else:  # config.year == "2010"
+            split_dp, images_dp, anns_dp = resource_dps
+
+            split_dp = Filter(split_dp, path_comparator("stem", config.split))
+            split_dp = LineReader(split_dp, decode=True, return_path=False)
+            split_dp = Mapper(split_dp, self._2010_split_key)
+
+            anns_dp = Mapper(anns_dp, self._2010_anns_key)
+
+        split_dp = hint_sharding(split_dp)
+        split_dp = hint_shuffling(split_dp)
+
+        dp = IterKeyZipper(
+            split_dp,
+            images_dp,
+            getitem(),
+            path_accessor("name"),
+            buffer_size=INFINITE_BUFFER_SIZE,
+        )
+        dp = IterKeyZipper(
+            dp,
+            anns_dp,
+            getitem(0),
+            buffer_size=INFINITE_BUFFER_SIZE,
+        )
+        return Mapper(dp, functools.partial(self._collate_and_decode_sample, year=config.year, decoder=decoder))
+
+    def _generate_categories(self, root: pathlib.Path) -> List[str]:
+        config = self.info.make_config(year="2011")
+        resources = self.resources(config)
+
+        dp = resources[0].load(root)
+        dp = Filter(dp, path_comparator("name", "classes.txt"))
+        dp = CSVDictParser(dp, fieldnames=("label", "category"), dialect="cub200")
+
+        return [row["category"].split(".")[1] for row in dp]