Source code for torchvision.datasets.gtsrb

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

[docs]class GTSRB(VisionDataset): """`German Traffic Sign Recognition Benchmark (GTSRB) <>`_ 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: 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)
[docs] def __getitem__(self, index: int) -> Tuple[Any, Any]: path, target = self._samples[index] sample ="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 = "" if self._split == "train": download_and_extract_archive( f"{base_url}", download_root=str(self._base_folder), md5="513f3c79a4c5141765e10e952eaa2478", ) else: download_and_extract_archive( f"{base_url}", download_root=str(self._base_folder), md5="c7e4e6327067d32654124b0fe9e82185", ) download_and_extract_archive( f"{base_url}", download_root=str(self._base_folder), md5="fe31e9c9270bbcd7b84b7f21a9d9d9e5", )


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources