abka03/stylsteer-vlm

1"""Abstract base class for style datasets."""2 3import logging4from abc import ABC, abstractmethod5from pathlib import Path6from typing import Any, Dict, List, Optional, Tuple7 8from PIL import Image9 10logger = logging.getLogger(__name__)11 12 13class StyleDataset(ABC):14    """Base class for all style captioning datasets.15 16    Each dataset provides:17    - Images for probing (train split) and evaluation (test split)18    - Style labels19    - Optionally: ground-truth styled captions (for GT metrics)20    """21 22    def __init__(23        self,24        data_dir: str,25        split: str = "test",26        n_images: Optional[int] = None,27        seed: int = 42,28    ):29        self.data_dir = Path(data_dir)30        self.split = split31        self.n_images = n_images32        self.seed = seed33        self._data: Optional[List[Dict[str, Any]]] = None34 35    @property36    @abstractmethod37    def track(self) -> str:38        """Track letter (A, B, C, D)."""39        ...40 41    @property42    @abstractmethod43    def styles(self) -> List[str]:44        """List of style names in this track."""45        ...46 47    @property48    @abstractmethod49    def has_ground_truth(self) -> bool:50        """Whether this dataset has ground-truth styled captions."""51        ...52 53    @abstractmethod54    def _load_data(self) -> List[Dict[str, Any]]:55        """Load raw data from disk.56 57        Returns list of dicts with at minimum:58            - "image_id": str or int59            - "image_path": str (absolute path to image)60            - "style": str (style name)61            - "caption_gt": Optional[List[str]] (ground-truth captions, if available)62        """63        ...64 65    @property66    def data(self) -> List[Dict[str, Any]]:67        """Lazy-loaded data."""68        if self._data is None:69            self._data = self._load_data()70        return self._data71 72    def get_images(self, style: str) -> List[Dict[str, Any]]:73        """Get all items for a given style."""74        items = [d for d in self.data if d["style"] == style]75        if self.n_images is not None:76            import random77            rng = random.Random(self.seed)78            items = rng.sample(items, min(self.n_images, len(items)))79        return items80 81    def load_image(self, image_path: str) -> Image.Image:82        """Load a PIL Image from path."""83        return Image.open(image_path).convert("RGB")84 85    def get_ground_truth(self, image_id: str, style: str) -> Optional[List[str]]:86        """Get ground-truth captions for an image+style pair."""87        if not self.has_ground_truth:88            return None89        items = [d for d in self.data if d["image_id"] == image_id and d["style"] == style]90        if not items:91            return None92        refs = []93        for item in items:94            if item.get("caption_gt"):95                refs.extend(item["caption_gt"] if isinstance(item["caption_gt"], list) else [item["caption_gt"]])96        return refs if refs else None97 98    def __len__(self) -> int:99        return len(self.data)100 101    def __repr__(self) -> str:102        return f"{self.__class__.__name__}(track={self.track}, split={self.split}, n={len(self)})"103