config/__init__.py

"""Configuration dataclasses for tactile_vae."""

from dataclasses import dataclass, field


@dataclass
class TactileVAEModelConfig:
    img_size: int = 128
    patch_size: int = 16
    in_chans: int = 3
    embed_dim: int = 256
    encoder_depth: int = 4
    encoder_heads: int = 8
    decoder_embed_dim: int = 192
    decoder_depth: int = 4
    decoder_heads: int = 8
    mlp_ratio: float = 4.0
    latent_dim: int = 128


@dataclass
class TactileVAELossConfig:
    beta: float = 1.0
    recon_type: str = "l1"
    ssim_weight: float = 0.0
    ssim_window_size: int = 11


@dataclass
class TactileVAETrainConfig:
    model: TactileVAEModelConfig = field(default_factory=TactileVAEModelConfig)
    loss: TactileVAELossConfig = field(default_factory=TactileVAELossConfig)