60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Callable, Optional, Union
|
||
|
|
|
||
|
|
from .folder import ImageFolder
|
||
|
|
from .utils import download_and_extract_archive
|
||
|
|
|
||
|
|
|
||
|
|
class EuroSAT(ImageFolder):
|
||
|
|
"""RGB version of the `EuroSAT <https://github.com/phelber/eurosat>`_ Dataset.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
root (str or ``pathlib.Path``): Root directory of dataset where ``root/eurosat`` exists.
|
||
|
|
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. Default is False.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
root: Union[str, Path],
|
||
|
|
transform: Optional[Callable] = None,
|
||
|
|
target_transform: Optional[Callable] = None,
|
||
|
|
download: bool = False,
|
||
|
|
) -> None:
|
||
|
|
self.root = os.path.expanduser(root)
|
||
|
|
self._base_folder = os.path.join(self.root, "eurosat")
|
||
|
|
self._data_folder = os.path.join(self._base_folder, "2750")
|
||
|
|
|
||
|
|
if download:
|
||
|
|
self.download()
|
||
|
|
|
||
|
|
if not self._check_exists():
|
||
|
|
raise RuntimeError("Dataset not found. You can use download=True to download it")
|
||
|
|
|
||
|
|
super().__init__(self._data_folder, transform=transform, target_transform=target_transform)
|
||
|
|
self.root = os.path.expanduser(root)
|
||
|
|
|
||
|
|
def __len__(self) -> int:
|
||
|
|
return len(self.samples)
|
||
|
|
|
||
|
|
def _check_exists(self) -> bool:
|
||
|
|
return os.path.exists(self._data_folder)
|
||
|
|
|
||
|
|
def download(self) -> None:
|
||
|
|
|
||
|
|
if self._check_exists():
|
||
|
|
return
|
||
|
|
|
||
|
|
os.makedirs(self._base_folder, exist_ok=True)
|
||
|
|
download_and_extract_archive(
|
||
|
|
"https://madm.dfki.de/files/sentinel/EuroSAT.zip",
|
||
|
|
download_root=self._base_folder,
|
||
|
|
md5="c8fa014336c82ac7804f0398fcb19387",
|
||
|
|
)
|