abka03/stylsteer-vlm

1"""Data loaders for StyleSteer-VLM."""2 3from src.data.base import StyleDataset4from src.data.senticap import SentiCapDataset5from src.data.flickrstyle10k import FlickrStyle10KDataset6from src.data.personality_caps import PersonalityCapsDataset7from src.data.coco import COCODataset8 9TRACK_DATASETS = {10    "A": SentiCapDataset,11    "B": FlickrStyle10KDataset,12    "C": PersonalityCapsDataset,13    "D": COCODataset,14}15 16 17def get_dataset(track: str, **kwargs) -> StyleDataset:18    """Get dataset by track letter."""19    if track not in TRACK_DATASETS:20        raise ValueError(f"Unknown track: {track}. Available: {list(TRACK_DATASETS.keys())}")21    return TRACK_DATASETS[track](**kwargs)22