config.py

import os
from pathlib import Path
try:
    import webvtt
    USE_WEBVTT = True
except ImportError:
    USE_WEBVTT = False
    print('webvtt-py not installed; using custom VTT parser')

# -- Project Paths ----------------------------------------------------------
# NOTE: Please update these paths to match your local machine's setup.
PROJECT_ROOT = Path(__file__).parent
CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
OUTPUT_DIR = PROJECT_ROOT / 'outputs'
BLEURT_CHECKPOINT_PATH = '/tmp/BLEURT-20'  # Multilingual BLEURT-20 (covers de_DE and zh_CN)

# -- Dataset switching ------------------------------------------------------
# Pick dataset via env var: DATASET=BOBSL (default) | PHOENIX | CSL. Each dataset is loaded from a drop-in BOBSL-style directory: 
# pose .npy + .vtt + subset2episode.json. PHOENIX/CSL synth is generated by data_synth/synthesize_streams.py.
#
# IMPORTANT: WIDTH/HEIGHT are the *native* canvas of each corpus; they are passed into poses/preprocessing.normalize_keypoints 
# (CoSign group-relative norm + global divide by W/H). We do NOT rescale Phoenix/CSL coordinates to BOBSL's 444 reference - 
# each language is trained from scratch on its own data, so its native canvas is the fair input distribution.
#
# mBART backbone: facebook/mbart-large-cc25 covers all 3 target languages (en_XX, de_DE, zh_CN);
DATASET = os.environ.get('DATASET', 'BOBSL').upper()
MBART_NAME = 'facebook/mbart-large-cc25'

# -- Backbone selection -----------------------------------------------------
# BACKBONE='cosign' (default) uses CoSign1s ST-GCN on 77 selected keypoints (group-norm).
# BACKBONE='mska' loads the multi-stream DSTA pose encoder vendored from MSKA, with weights from MSKA SLT release (gloss/CTC heads stripped at load-time).
# MSKA_CKPT auto-resolves to the per-dataset checkpoint; override via env var if needed.
BACKBONE = os.environ.get('BACKBONE', 'cosign').lower()
if BACKBONE not in ('cosign', 'mska'): raise ValueError(f"Unknown BACKBONE={BACKBONE!r}; expected 'cosign' or 'mska'")
if DATASET == 'BOBSL':
    DATA_ROOT = Path('data/BOBSL')
    POSE_ROOT = DATA_ROOT / 'bobsl_dwpose' # Directory with pose .npy files
    SUBSET_JSON = DATA_ROOT / 'original_data/metadata/subset2episode.json' # subset2episode.json for Train/val/test splits
    VTT_DIR = DATA_ROOT / 'automatic_annotations/signing_aligned_subtitles/auto_sat_aligned' # Directory with .vtt files
    TGT_LANG = 'en_XX'
    TRIMMED_MBART_DIR = TRIMMED_TOKENIZER_DIR = 'captioners/trimmed_mbart_bobsl'
    WIDTH, HEIGHT, FPS = 444, 444, 12.5  # Downsampled FPS from original 25fps
elif DATASET == 'PHOENIX':
    DATA_ROOT = Path('data/synth/phoenix')
    POSE_ROOT = DATA_ROOT / 'poses'
    SUBSET_JSON = DATA_ROOT / 'subset2episode.json'
    VTT_DIR = DATA_ROOT / 'vtt'
    TGT_LANG = 'de_DE'
    TRIMMED_MBART_DIR = TRIMMED_TOKENIZER_DIR = 'captioners/trimmed_mbart_phoenix'
    WIDTH, HEIGHT, FPS = 210, 260, 12.5  # Downsampled FPS from original 25fps
    # Legacy big-pickle layout (one {.train,.dev,.test} pickle, each a dict of {name: {'keypoint': (T,133,3), 'text': ...}}).
    # `pickle_prefix` present -> synthesize_streams.py uses the legacy loader path (NOT Uni-Sign per-sample).
    SYNTH_META = {'src_fps': 25.0, 'src_w': WIDTH, 'src_h': HEIGHT, 'pickle_dir': Path('data/Phoenix-2014T'),
                  'pickle_prefix': 'Phoenix-2014T', 'splits': {'train': 'train', 'val': 'dev', 'test': 'test'}}
elif DATASET == 'CSL':
    DATA_ROOT = Path('data/synth/csl')
    POSE_ROOT = DATA_ROOT / 'poses'
    SUBSET_JSON = DATA_ROOT / 'subset2episode.json'
    VTT_DIR = DATA_ROOT / 'vtt'
    TGT_LANG = 'zh_CN'
    TRIMMED_MBART_DIR = TRIMMED_TOKENIZER_DIR = 'captioners/trimmed_mbart_csl'
    WIDTH, HEIGHT, FPS = 512, 512, 15   # Padded reference for CSL-Daily
    # Per-sample layout: pose_dir/<name>.pkl (RTMPose normalized + scores + w_h, with start/end). `pickle_prefix` absent.
    SYNTH_META = {'src_fps': 30.0, 'src_w': WIDTH, 'src_h': HEIGHT,
                  'pickle_dir': Path('data/CSL-Daily'), 'pose_dir_name': 'poses', 'label_prefix': 'labels',
                  'splits': {'train': 'train', 'val': 'dev', 'test': 'test'}}
elif DATASET == 'H2S':
    # We use the SAME BOBSL-style stream synthesis as CSL/PHOENIX (signer-pure pools, BOBSL empirical pauses, Hermite C1 bridges) 
    # -- Real CSV inter-sentence gaps were dropped because >10s instructional dead-time produces too many empty training windows.
    DATA_ROOT = Path('data/synth/h2s')
    POSE_ROOT = DATA_ROOT / 'poses'
    SUBSET_JSON = DATA_ROOT / 'subset2episode.json'
    VTT_DIR = DATA_ROOT / 'vtt'
    TGT_LANG = 'en_XX'
    TRIMMED_MBART_DIR = TRIMMED_TOKENIZER_DIR = 'captioners/trimmed_mbart_h2s'
    WIDTH, HEIGHT, FPS = 1280, 720, 15  # Global rescale canvas; per-sample w_h varies and is normalized away.
    # Per-sample layout: pose_dir/<name>.pkl (RTMPose normalized + scores + w_h, with start/end). `pickle_prefix` absent.
    SYNTH_META = {'src_fps': 30.0, 'src_w': WIDTH, 'src_h': HEIGHT,
                  'pickle_dir': Path('data/How2Sign'), 'pose_dir_name': 'poses', 'label_prefix': 'labels',
                  'splits': {'train': 'train', 'test': 'test'}}  # no 'val' key by design
else: raise ValueError(f"Unknown DATASET={DATASET!r}; expected one of BOBSL/PHOENIX/CSL/H2S")

# Optional data-root override for drop-in-layout variants of the active dataset (e.g. the clip-level baseline). Everything else -- tokenizer, language, fps, 
# canvas -- stays as configured by DATASET above.
if os.environ.get('SLT_DATA_ROOT'):
    DATA_ROOT = Path(os.environ['SLT_DATA_ROOT'])
    POSE_ROOT = DATA_ROOT / 'poses'
    SUBSET_JSON = DATA_ROOT / 'subset2episode.json'
    VTT_DIR = DATA_ROOT / 'vtt'
if os.environ.get('SLT_FPS'): FPS = float(os.environ['SLT_FPS']) # e.g. SLT_FPS=25 for the full-frame-rate clip ablation

# MSKA checkpoint auto-resolution per dataset. Override with MSKA_CKPT env var.
_MSKA_CKPT_DEFAULTS = {'PHOENIX': 'checkpoints/mska_phoenix.pth', 'CSL':     'checkpoints/mska_csl.pth'}
MSKA_CKPT = os.environ.get('MSKA_CKPT', _MSKA_CKPT_DEFAULTS.get(DATASET, ''))

# -- Dataset and Dataloader Configuration -----------------------------------
MIN_SUB_DURATION = 1.0  # From LiTFiC, seconds
MAX_SUB_DURATION = 20.0 # From LiTFiC, seconds
WINDOW_DURATION_SECONDS = 15 # As per https://aclanthology.org/2025.acl-srw.93.pdf

# -- Pose Preprocessing (CoSign Inspired) -----------------------------------
# Define keypoint groups based on COCO-WholeBody (133 points: body 0-16, left foot 17-19, right foot 20-22, face 23-90, left hand 91-111, right hand 112-132)
# Upper body (9): nose(0), left eye(1), right eye(2), left shoulder(5), right shoulder(6), left elbow(7), right elbow(8), left wrist(9), right wrist(10)
# Mouth (8): inner lips approx 60-67 in face (face starts at 23, so 23+60-67 = 83-90, but adjust to 8)
# Face lower/cheek (18): contour approx 23+0 to 23+16 (17 points), plus nose
BODY_IDS = [0, 1, 2, 5, 6, 7, 8, 9, 10]  # 9 points
LEFT_HAND_IDS = list(range(91, 112))  # 21 points
RIGHT_HAND_IDS = list(range(112, 133))  # 21 points
MOUTH_IDS = list(range(83, 91))  # Inner mouth 8 points
FACE_IDS = list(range(23, 40)) + [53]  # First 18 as cheek/lower approx
ALL_SELECTED_IDS = BODY_IDS + LEFT_HAND_IDS + RIGHT_HAND_IDS + MOUTH_IDS + FACE_IDS  # Total 9+21+21+8+18=77
CONF_THRESHOLD = float(os.environ.get('CONF_THRESHOLD', 0.5))  # From supp: Keypoints with conf > 0.5 considered valid
NUM_KEYPOINTS = len(ALL_SELECTED_IDS)
KPS_MODULES = {
    'body': {'kps_ids': BODY_IDS, 'kps_rel_range': (0, 9)},
    'left_hand': {'kps_ids': LEFT_HAND_IDS, 'kps_rel_range': (9, 30)},
    'right_hand': {'kps_ids': RIGHT_HAND_IDS, 'kps_rel_range': (30, 51)},
    'mouth': {'kps_ids': MOUTH_IDS, 'kps_rel_range': (51, 59)},
    'face': {'kps_ids': FACE_IDS, 'kps_rel_range': (59, 77)},
}