thienphuc12339/SignLanguage

1from pathlib import Path2from typing import Any3from dataclasses import dataclass, field4from utils import MODELS, VIDEO_EXTENSIONS5 6 7@dataclass8class TransformConfig:9    # RGB specific10    horizontal_flip_prob: float = 0.511    aug_type: str = "augmix"12    aug_paras: dict = field(13        default_factory=lambda: {14            "magnitude": 3,15            "alpha": 1.0,16            "width": 5,17            "depth": -1,18        }19    )20    sample_rate: int = 421 22    # Pose specific23    normalization: bool = True24 25    # SL-GCN, DSTA-SLR specific26    random_choose: bool = False27    random_shift: bool = False28    random_move: bool = False29    random_mirror: bool = False30    random_mirror_p: float = 0.531    bone_stream: bool = False32    motion_stream: bool = False33 34    # SPOTER specific35    augmentation: bool = True36    aug_prob: float = 0.537    noise: bool = True38 39    def __post_init__(self):40        assert self.aug_type in ["augmix", "mixup"], \41            "Only AugMix and MixUp are supported for now"42 43 44@dataclass45class DataConfig:46    dataset: str = "vsl"47    modality: str = "rgb"48    subset: str = None49    data_dir: str = "data/processed/vsl"50    transform: Any = None51    fps: int = 3052    debug: bool = False53    # transform: TransformConfig = TransformConfig()54    transform: TransformConfig = field(default_factory=TransformConfig)55 56 57    def __post_init__(self):58        assert self.dataset in ["vsl_98", "vsl_400"], \59            "Only VSL dataset is supported for now"60        assert self.modality in ["rgb", "pose"], \61            "Only RGB and Pose modalities are supported for now"62 63 64@dataclass65class ModelConfig:66    arch: str = "sl_gcn"67    pretrained: str = "vsltranslation/sl_gcn_joint_v3_0"68    num_frozen_layers: int = 069    ignored_weights: list = field(default_factory=lambda: [])70    num_frames: int = 1671 72    # SL-GCN specific73    num_points: int = 2774    groups: int = 875    block_size: int = 4176    in_channels: int = 377    labeling_mode: str = "spatial"78    is_vector: bool = False79 80    # DSTA-SLR specific81    graph: str = "wlasl"82    inner_dim: int = 6483    drop_layers: int = 284    depth: int = 485    s_num_heads: int = 186    window_size: int = 12087 88    # SPOTER specific89    hidden_dim: int = 10890 91    def __post_init__(self):92        assert self.arch in MODELS, f"Model {self.arch} is not supported"93 94 95@dataclass96class TrainingConfig:97    output_dir: str = "experiments"98    remove_unused_columns: bool = False99    do_train: bool = True100    use_cpu: bool = False101 102    eval_strategy: str = "epoch"103    logging_strategy: str = "epoch"104    save_strategy: str = "epoch"105    logging_steps: int = 1106    save_steps: int = 1107    eval_steps: int = 1108 109    learning_rate: float = 5e-5110    weight_decay: float = 0111    adam_beta1: float = 0.9112    adam_beta2: float = 0.999113    adam_epsilon: float = 1e-8114    warmup_ratio: float = 0.1115 116    num_train_epochs: int = 10117    per_device_train_batch_size: int = 8118    per_device_eval_batch_size: int = 8119    dataloader_num_workers: int = 0120 121    load_best_model_at_end: bool = True122    metric_for_best_model: str = "accuracy"123    resume_from_checkpoint: str = None124 125    run_name: str = "swin3d"126    report_to: str = None127    push_to_hub: bool = False128    hub_model_id: str = None129    hub_strategy: str = "checkpoint"130    hub_private_repo: bool = True131 132    def __post_init__(self):133        self.output_dir = Path(self.output_dir)134        if str(self.output_dir) == "experiments":135            self.output_dir = self.output_dir / self.run_name136        self.output_dir.mkdir(parents=True, exist_ok=True)137 138        if self.hub_model_id is not None:139            self.push_to_hub = True140            if len(self.hub_model_id.split("/")) == 1:141                self.hub_model_id = f"{self.hub_model_id}/{self.run_name}"142 143 144@dataclass145class InferenceConfig:146    source: str = "webcam"147    output_dir: str = "demo"148    use_onnx: bool = False149    device: str = "cpu"150    cache_dir: str = "models/huggingface"151 152    visualize: bool = False153    show_skeleton: bool = False154 155    visibility: float = 0.5156    angle_threshold: int = 140157    min_num_up_frames: int = 10158    min_num_down_frames: int = 10159    delay: int = 400160 161    top_k: int = 3162    # SL-GCN, DSTA-SLR specific163    bone_stream: bool = False164    motion_stream: bool = False165 166    def __post_init__(self):167        self.source = Path(self.source)168        assert any((169            str(self.source) == "webcam",170            (self.source.exists() and str(self.source).endswith(VIDEO_EXTENSIONS))171        )), \172            f"Only Webcam and Video sources are supported for now (got {self.source})"173        self.output_dir = Path(self.output_dir)174        self.output_dir.mkdir(parents=True, exist_ok=True)175