# Wraping for the EuroSAT dataset

from typing import Callable, Optional

from torchvision.datasets import EuroSAT as EuroSAT_tv
from torchvision.transforms import transforms


class EuroSAT(EuroSAT_tv):

    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:

        super().__init__(
            root,
            transforms.ToTensor() if transform is None else transform,
            target_transform, download)

        self.classes_names = self.classes

        self.mean, self.std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
        if isinstance(self.transform, transforms.Compose):
            self.transform.transforms.append(
                transforms.Normalize(self.mean, self.std))
