# Generated by configen, do not edit.
# See https://github.com/facebookresearch/hydra/tree/master/tools/configen
# fmt: off
# isort:skip_file
# flake8: noqa

from dataclasses import dataclass, field
from conduit.data.datasets.vision.camelyon17 import Camelyon17Attr
from conduit.data.datasets.vision.camelyon17 import Camelyon17SplitScheme
from conduit.data.datasets.vision.celeba import CelebAttr
from conduit.data.datasets.vision.nico import NicoSuperclass
from omegaconf import MISSING
from pathlib import Path
from ranzen.torch.data import TrainingMode
from typing import Any
from typing import List
from typing import Optional
from typing import Union


@dataclass
class CelebADataModuleConf:
    _target_: str = "conduit.data.datamodules.vision.CelebADataModule"
    root: Union[Path, str] = MISSING
    train_batch_size: int = 64
    eval_batch_size: Optional[int] = None
    val_prop: float = 0.2
    test_prop: float = 0.2
    num_workers: int = 0
    seed: int = 47
    persist_workers: bool = False
    pin_memory: bool = True
    stratified_sampling: bool = False
    instance_weighting: bool = False
    training_mode: TrainingMode = TrainingMode.epoch
    train_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    test_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    image_size: int = 224
    superclass: CelebAttr = CelebAttr.SMILING
    subclass: CelebAttr = CelebAttr.MALE
    use_predefined_splits: bool = False


@dataclass
class ColoredMNISTDataModuleConf:
    _target_: str = "conduit.data.datamodules.vision.ColoredMNISTDataModule"
    root: Union[Path, str] = MISSING
    train_batch_size: int = 64
    eval_batch_size: Optional[int] = None
    val_prop: float = 0.2
    test_prop: float = 0.2
    num_workers: int = 0
    seed: int = 47
    persist_workers: bool = False
    pin_memory: bool = True
    stratified_sampling: bool = False
    instance_weighting: bool = False
    training_mode: TrainingMode = TrainingMode.epoch
    train_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    test_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    image_size: int = 32
    use_predefined_splits: bool = False
    label_map: Any = None  # Optional[Dict[int, int]]
    colors: Optional[List[int]] = None
    num_colors: int = 10
    scale: float = 0.2
    correlation: float = 1.0
    binarize: bool = False
    greyscale: bool = False
    background: bool = False
    black: bool = True


@dataclass
class NICODataModuleConf:
    _target_: str = "conduit.data.datamodules.vision.NICODataModule"
    root: Union[Path, str] = MISSING
    train_batch_size: int = 64
    eval_batch_size: Optional[int] = None
    val_prop: float = 0.2
    test_prop: float = 0.2
    num_workers: int = 0
    seed: int = 47
    persist_workers: bool = False
    pin_memory: bool = True
    stratified_sampling: bool = False
    instance_weighting: bool = False
    training_mode: TrainingMode = TrainingMode.epoch
    train_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    test_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    image_size: int = 224
    class_train_props: Optional[dict] = None
    superclass: NicoSuperclass = NicoSuperclass.ANIMALS


@dataclass
class WaterbirdsDataModuleConf:
    _target_: str = "conduit.data.datamodules.vision.WaterbirdsDataModule"
    root: Union[Path, str] = MISSING
    train_batch_size: int = 64
    eval_batch_size: Optional[int] = None
    val_prop: float = 0.2
    test_prop: float = 0.2
    num_workers: int = 0
    seed: int = 47
    persist_workers: bool = False
    pin_memory: bool = True
    stratified_sampling: bool = False
    instance_weighting: bool = False
    training_mode: TrainingMode = TrainingMode.epoch
    train_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    test_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    image_size: int = 224
    use_predefined_splits: bool = False


@dataclass
class Camelyon17DataModuleConf:
    _target_: str = "conduit.data.datamodules.vision.Camelyon17DataModule"
    root: Union[Path, str] = MISSING
    train_batch_size: int = 64
    eval_batch_size: Optional[int] = None
    val_prop: float = 0.2
    test_prop: float = 0.2
    num_workers: int = 0
    seed: int = 47
    persist_workers: bool = False
    pin_memory: bool = True
    stratified_sampling: bool = False
    instance_weighting: bool = False
    training_mode: TrainingMode = TrainingMode.epoch
    train_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    test_tf: Any = None  # Optional[Union[BasicTransform, Callable[[Image], Any], Compose]]
    image_size: int = 96
    superclass: Camelyon17Attr = Camelyon17Attr.TUMOR
    subclass: Camelyon17Attr = Camelyon17Attr.CENTER
    use_predefined_splits: bool = False
    split_scheme: Camelyon17SplitScheme = Camelyon17SplitScheme.OFFICIAL
