104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
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",
|
|
)
|