OxfordIIITPet dataset

docs/source/datasets.rst

@@ -58,6 +58,7 @@ You can also create your own datasets using the provided :ref:`base classes <bas
     LSUN
     MNIST
     Omniglot
+    OxfordIIITPet
     PhotoTour
     Places365
     QMNIST

test/test_datasets.py

@@ -1,5 +1,6 @@
 import bz2
 import contextlib
+import csv
 import io
 import itertools
 import json
@@ -2168,5 +2169,64 @@ def inject_fake_data(self, tmpdir, config):
         return num_sequences * (num_examples_per_sequence - 1)
 
 
+class OxfordIIITPetTestCase(datasets_utils.DatasetTestCase):
+    DATASET_CLASS = datasets.OxfordIIITPet
+    FEATURE_TYPES = (PIL.Image.Image, (int, PIL.Image.Image, tuple, type(None)))
+    ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(
+        split=("trainval", "test"),
+        target_type=("category", "segmentation", ("category", "segmentation")),
+    )
+
+    def inject_fake_data(self, tmpdir, config):
+        base_folder = os.path.join(tmpdir, "oxford-iiit-pet")
+
+        classification_anns_meta = (
+            dict(cls="Abyssinian", label=0, species="cat"),
+            dict(cls="Keeshond", label=18, species="dog"),
+            dict(cls="Yorkshire Terrier", label=37, species="dog"),
+        )
+        split_and_classification_anns = [
+            self._meta_to_split_and_classification_ann(meta, idx)
+            for meta, idx in itertools.product(classification_anns_meta, (1, 2, 10))
+        ]
+        image_ids, *_ = zip(*split_and_classification_anns)
+
+        image_files = datasets_utils.create_image_folder(
+            base_folder, "images", file_name_fn=lambda idx: f"{image_ids[idx]}.jpg", num_examples=len(image_ids)
+        )
+
+        anns_folder = os.path.join(base_folder, "annotations")
+        os.makedirs(anns_folder)
+        split_and_classification_anns_in_split = random.choices(split_and_classification_anns, k=len(image_ids) // 2)
+        with open(os.path.join(anns_folder, f"{config['split']}.txt"), "w", newline="") as file:
+            writer = csv.writer(file, delimiter=" ")
+            for split_and_classification_ann in split_and_classification_anns_in_split:
+                writer.writerow(split_and_classification_ann)
+
+        segmentation_files = datasets_utils.create_image_folder(
+            anns_folder, "trimaps", file_name_fn=lambda idx: f"{image_ids[idx]}.png", num_examples=len(image_ids)
+        )
+
+        # The dataset has some rouge files
+        for path in image_files[:2]:
+            path.with_suffix(".mat").touch()
+        for path in segmentation_files:
+            path.with_name(f".{path.name}").touch()
+
+        return len(split_and_classification_anns_in_split)
+
+    def _meta_to_split_and_classification_ann(self, meta, idx):
+        image_id = "_".join(
+            [
+                *[(str.title if meta["species"] == "cat" else str.lower)(part) for part in meta["cls"].split()],
+                str(idx),
+            ]
+        )
+        class_id = str(meta["label"] + 1)
+        species = "1" if meta["species"] == "cat" else "2"
+        breed_id = "-1"
+        return (image_id, class_id, species, breed_id)
+
+
 if __name__ == "__main__":
     unittest.main()

torchvision/datasets/__init__.py

@@ -16,6 +16,7 @@
 from .lsun import LSUN, LSUNClass
 from .mnist import MNIST, EMNIST, FashionMNIST, KMNIST, QMNIST
 from .omniglot import Omniglot
+from .oxford_iiit_pet import OxfordIIITPet
 from .phototour import PhotoTour
 from .places365 import Places365
 from .sbd import SBDataset
@@ -77,4 +78,5 @@
     "FlyingChairs",
     "FlyingThings3D",
     "HD1K",
+    "OxfordIIITPet",
 )

torchvision/datasets/oxford_iiit_pet.py

@@ -0,0 +1,115 @@
+import os
+import os.path
+from typing import Any, Callable, Optional, Union, Tuple
+from typing import Sequence
+
+from PIL import Image
+
+from .utils import download_and_extract_archive, verify_str_arg
+from .vision import VisionDataset
+
+
+class OxfordIIITPet(VisionDataset):
+    _RESOURCES = (
+        ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz", "5c4f3ee8e5d25df40f4fd59a7f44e54c"),
+        ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz", "95a8c909bbe2e81eed6a22bccdf3f68f"),
+    )
+    _TARGET_TYPES = ("category", "segmentation")
+
+    def __init__(
+        self,
+        root: str,
+        split: str = "trainval",
+        target_type: Union[Sequence[str], str] = "category",
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        download: bool = True,
+    ):
+        super().__init__(root, transforms=transforms, transform=transform, target_transform=target_transform)
+        self._split = verify_str_arg(split, "split", ("trainval", "test"))
+        if isinstance(target_type, str):
+            target_type = [target_type]
+        self._target_type = [verify_str_arg(t, "target_type", self._TARGET_TYPES) for t in target_type]
+
+        if download:
+            self._download()
+
+        if not self._check_exists():
+            raise RuntimeError("Dataset not found. You can use download=True to download it")
+
+        image_ids = []
+        self._labels = []
+        with open(os.path.join(self._anns_folder, f"{self._split}.txt")) as file:
+            for line in file:
+                image_id, label, *_ = line.strip().split()
+                image_ids.append(image_id)
+                self._labels.append(int(label) - 1)
+
+        self.classes = [
+            " ".join(part.title() for part in raw_cls.split("_"))
+            for raw_cls, _ in sorted(
+                {(image_id.rsplit("_", 1)[0], label) for image_id, label in zip(image_ids, self._labels)},
+                key=lambda image_id_and_label: image_id_and_label[1],
+            )
+        ]
+        self.class_to_idx = dict(zip(self.classes, range(len(self.classes))))
+
+        self._images = [os.path.join(self._images_folder, f"{image_id}.jpg") for image_id in image_ids]
+        self._segmentations = [os.path.join(self._segmentations_folder, f"{image_id}.png") for image_id in image_ids]
+
+    def __len__(self) -> int:
+        return len(self._images)
+
+    def __getitem__(self, idx: int) -> Tuple[Any, Any]:
+        image = Image.open(self._images[idx]).convert("RGB")
+
+        target = []
+        for t in self._target_type:
+            if t == "category":
+                target.append(self._labels[idx])
+            else:  # t == "segmentation"
+                target.append(Image.open(self._segmentations[idx]))
+
+        target: Any
+        if not target:
+            target = None
+        elif len(target) == 1:
+            target = target[0]
+        else:
+            target = tuple(target)
+
+        if self.transforms:
+            image, target = self.transforms(image, target)
+
+        return image, target
+
+    @property
+    def _base_folder(self):
+        return os.path.join(self.root, "oxford-iiit-pet")
+
+    @property
+    def _images_folder(self) -> str:
+        return os.path.join(self._base_folder, "images")
+
+    @property
+    def _anns_folder(self) -> str:
+        return os.path.join(self._base_folder, "annotations")
+
+    @property
+    def _segmentations_folder(self) -> str:
+        return os.path.join(self._anns_folder, "trimaps")
+
+    def _check_exists(self) -> bool:
+        for folder in (self._images_folder, self._anns_folder):
+            if not (os.path.exists(folder) and os.path.isdir(folder)):
+                return False
+        else:
+            return True
+
+    def _download(self) -> None:
+        if self._check_exists():
+            return
+
+        for url, md5 in self._RESOURCES:
+            download_and_extract_archive(url, download_root=self._base_folder, md5=md5)

torchvision/prototype/datasets/_builtin/__init__.py

@@ -4,6 +4,7 @@
 from .coco import Coco
 from .imagenet import ImageNet
 from .mnist import MNIST, FashionMNIST, KMNIST, EMNIST, QMNIST
+from .oxford_iiit_pet import OxfordIITPet
 from .sbd import SBD
 from .semeion import SEMEION
 from .voc import VOC

torchvision/prototype/datasets/_builtin/oxford-iiit-pet.categories

@@ -0,0 +1,37 @@
+Abyssinian
+American Bulldog
+American Pit Bull Terrier
+Basset Hound
+Beagle
+Bengal
+Birman
+Bombay
+Boxer
+British Shorthair
+Chihuahua
+Egyptian Mau
+English Cocker Spaniel
+English Setter
+German Shorthaired
+Great Pyrenees
+Havanese
+Japanese Chin
+Keeshond
+Leonberger
+Maine Coon
+Miniature Pinscher
+Newfoundland
+Persian
+Pomeranian
+Pug
+Ragdoll
+Russian Blue
+Saint Bernard
+Samoyed
+Scottish Terrier
+Shiba Inu
+Siamese
+Sphynx
+Staffordshire Bull Terrier
+Wheaten Terrier
+Yorkshire Terrier