| import csv | |
| import pathlib | |
| from typing import Any, Callable, Optional, Tuple, Union | |
| import PIL | |
| from .folder import make_dataset | |
| from .utils import download_and_extract_archive, verify_str_arg | |
| from .vision import VisionDataset | |
| class GTSRB(VisionDataset): | |
| """`German Traffic Sign Recognition Benchmark (GTSRB) <https://benchmark.ini.rub.de/>`_ Dataset. | |
| Args: | |
| root (str or ``pathlib.Path``): Root directory of the dataset. | |
| split (string, optional): The dataset split, supports ``"train"`` (default), or ``"test"``. | |
| transform (callable, optional): A function/transform that takes in a PIL image 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 in root directory. If dataset is already downloaded, it is not | |
| downloaded again. | |
| """ | |
| def __init__( | |
| self, | |
| root: Union[str, pathlib.Path], | |
| split: str = "train", | |
| transform: Optional[Callable] = None, | |
| target_transform: Optional[Callable] = None, | |
| download: bool = False, | |
| ) -> None: | |
| super().__init__(root, transform=transform, target_transform=target_transform) | |
| self._split = verify_str_arg(split, "split", ("train", "test")) | |
| self._base_folder = pathlib.Path(root) / "gtsrb" | |
| self._target_folder = ( | |
| self._base_folder / "GTSRB" / ("Training" if self._split == "train" else "Final_Test/Images") | |
| ) | |
| if download: | |
| self.download() | |
| if not self._check_exists(): | |
| raise RuntimeError("Dataset not found. You can use download=True to download it") | |
| if self._split == "train": | |
| samples = make_dataset(str(self._target_folder), extensions=(".ppm",)) | |
| else: | |
| with open(self._base_folder / "GT-final_test.csv") as csv_file: | |
| samples = [ | |
| (str(self._target_folder / row["Filename"]), int(row["ClassId"])) | |
| for row in csv.DictReader(csv_file, delimiter=";", skipinitialspace=True) | |
| ] | |
| self._samples = samples | |
| self.transform = transform | |
| self.target_transform = target_transform | |
| def __len__(self) -> int: | |
| return len(self._samples) | |
| def __getitem__(self, index: int) -> Tuple[Any, Any]: | |
| path, target = self._samples[index] | |
| sample = PIL.Image.open(path).convert("RGB") | |
| if self.transform is not None: | |
| sample = self.transform(sample) | |
| if self.target_transform is not None: | |
| target = self.target_transform(target) | |
| return sample, target | |
| def _check_exists(self) -> bool: | |
| return self._target_folder.is_dir() | |
| def download(self) -> None: | |
| if self._check_exists(): | |
| return | |
| base_url = "https://sid.erda.dk/public/archives/daaeac0d7ce1152aea9b61d9f1e19370/" | |
| if self._split == "train": | |
| download_and_extract_archive( | |
| f"{base_url}GTSRB-Training_fixed.zip", | |
| download_root=str(self._base_folder), | |
| md5="513f3c79a4c5141765e10e952eaa2478", | |
| ) | |
| else: | |
| download_and_extract_archive( | |
| f"{base_url}GTSRB_Final_Test_Images.zip", | |
| download_root=str(self._base_folder), | |
| md5="c7e4e6327067d32654124b0fe9e82185", | |
| ) | |
| download_and_extract_archive( | |
| f"{base_url}GTSRB_Final_Test_GT.zip", | |
| download_root=str(self._base_folder), | |
| md5="fe31e9c9270bbcd7b84b7f21a9d9d9e5", | |
| ) | |