| from pathlib import Path | |
| from typing import Any, Callable, Optional, Union | |
| from .folder import default_loader, ImageFolder | |
| from .utils import download_and_extract_archive, verify_str_arg | |
| class Country211(ImageFolder): | |
| """`The Country211 Data Set <https://github.com/openai/CLIP/blob/main/data/country211.md>`_ from OpenAI. | |
| This dataset was built by filtering the images from the YFCC100m dataset | |
| that have GPS coordinate corresponding to a ISO-3166 country code. The | |
| dataset is balanced by sampling 150 train images, 50 validation images, and | |
| 100 test images for each country. | |
| Args: | |
| root (str or ``pathlib.Path``): Root directory of the dataset. | |
| split (string, optional): The dataset split, supports ``"train"`` (default), ``"valid"`` and ``"test"``. | |
| transform (callable, optional): A function/transform that takes in a PIL image or torch.Tensor, depends on the given loader, | |
| and returns a transformed version. E.g, ``transforms.RandomCrop`` | |
| target_transform (callable, optional): A function/transform that takes in the target and transforms it. | |
| download (bool, optional): If True, downloads the dataset from the internet and puts it into | |
| ``root/country211/``. If dataset is already downloaded, it is not downloaded again. | |
| loader (callable, optional): A function to load an image given its path. | |
| By default, it uses PIL as its image loader, but users could also pass in | |
| ``torchvision.io.decode_image`` for decoding image data into tensors directly. | |
| """ | |
| _URL = "https://openaipublic.azureedge.net/clip/data/country211.tgz" | |
| _MD5 = "84988d7644798601126c29e9877aab6a" | |
| def __init__( | |
| self, | |
| root: Union[str, Path], | |
| split: str = "train", | |
| transform: Optional[Callable] = None, | |
| target_transform: Optional[Callable] = None, | |
| download: bool = False, | |
| loader: Callable[[str], Any] = default_loader, | |
| ) -> None: | |
| self._split = verify_str_arg(split, "split", ("train", "valid", "test")) | |
| root = Path(root).expanduser() | |
| self.root = str(root) | |
| self._base_folder = root / "country211" | |
| if download: | |
| self._download() | |
| if not self._check_exists(): | |
| raise RuntimeError("Dataset not found. You can use download=True to download it") | |
| super().__init__( | |
| str(self._base_folder / self._split), | |
| transform=transform, | |
| target_transform=target_transform, | |
| loader=loader, | |
| ) | |
| self.root = str(root) | |
| def _check_exists(self) -> bool: | |
| return self._base_folder.exists() and self._base_folder.is_dir() | |
| def _download(self) -> None: | |
| if self._check_exists(): | |
| return | |
| download_and_extract_archive(self._URL, download_root=self.root, md5=self._MD5) | |