abka03/stylsteer-vlm

1"""Track D — COCO val2017 dataset loader.2 3Used for custom rhetorical styles (poetic, scientific, narrative, factual,4natural, fictional). No ground-truth styled captions — LLM judge only.5"""6 7import json8import logging9import random10from pathlib import Path11from typing import Any, Dict, List, Optional12 13from src.data.base import StyleDataset14 15logger = logging.getLogger(__name__)16 17 18class COCODataset(StyleDataset):19    """COCO val2017 for custom rhetorical tracks (Track D).20 21    Expected directory structure:22        COCO_IMAGE_DIR/   (from .env, e.g. data/coco_val2017/images/)23            ├── 000000000139.jpg24            └── ...25        COCO_ANNOT_FILE/  (from .env)26            └── captions_val2017.json27 28    No ground-truth styled captions. Factual COCO captions used as neutral reference only.29    """30 31    RHETORICAL_STYLES = ["poetic", "scientific", "narrative", "factual", "natural", "fictional"]32 33    def __init__(34        self,35        data_dir: str = "",36        image_dir: str = "",37        annot_file: str = "",38        **kwargs,39    ):40        self.image_dir = Path(image_dir) if image_dir else None41        self.annot_file = Path(annot_file) if annot_file else None42        super().__init__(data_dir=data_dir, **kwargs)43 44    @property45    def track(self) -> str:46        return "D"47 48    @property49    def styles(self) -> List[str]:50        return self.RHETORICAL_STYLES51 52    @property53    def has_ground_truth(self) -> bool:54        return False55 56    def _load_data(self) -> List[Dict[str, Any]]:57        # Resolve image directory58        img_dir = self.image_dir or self.data_dir / "images"59        annot = self.annot_file or self.data_dir / "annotations" / "captions_val2017.json"60 61        if not Path(img_dir).exists():62            logger.warning(f"COCO images not found at {img_dir}. Using mock data.")63            return self._mock_data()64 65        # Load annotation file to get image list66        image_ids = []67        if Path(annot).exists():68            with open(annot) as f:69                coco_data = json.load(f)70            image_ids = [img["id"] for img in coco_data.get("images", [])]71        else:72            # Fall back to listing image files73            image_ids = [74                int(p.stem) for p in Path(img_dir).glob("*.jpg")75            ]76 77        # Deterministic subset78        rng = random.Random(self.seed)79        n = self.n_images or 50080        if len(image_ids) > n:81            image_ids = rng.sample(image_ids, n)82        image_ids.sort()83 84        # Create entries — one per image per style85        data = []86        for img_id in image_ids:87            filename = f"{img_id:012d}.jpg"88            image_path = str(Path(img_dir) / filename)89            for style in self.styles:90                data.append({91                    "image_id": str(img_id),92                    "image_path": image_path,93                    "style": style,94                    "caption_gt": None,  # No ground truth for Track D95                })96 97        logger.info(f"COCO Track D: {len(image_ids)} images × {len(self.styles)} styles = {len(data)} items")98        return data99 100    def get_images(self, style: str) -> List[Dict[str, Any]]:101        """Override to return unique images (not per-style duplicates)."""102        items = [d for d in self.data if d["style"] == style]103        # Already de-duplicated at image level since each image has one entry per style104        return items105 106    def _mock_data(self) -> List[Dict[str, Any]]:107        data = []108        for i in range(max(self.n_images or 5, 5)):109            for style in self.styles:110                data.append({111                    "image_id": f"mock_{i}",112                    "image_path": f"mock_image_{i}.jpg",113                    "style": style,114                    "caption_gt": None,115                })116        return data117