training/prepare_dataset.py

"""
Dataset preparation for LoRA training.

For each heritage style:
  1. Source 30-50 high-quality reference images (public domain museum archives,
     Wikipedia Commons, or AI-generated seeds)
  2. Resize to 1024×1024
  3. Generate captions using the style's cultural keywords + a generic description
  4. Output to assets/datasets/<style>/image_NN.jpg + image_Nn.txt (caption)

The captions are critical — they teach the LoRA to associate the visual style
with the cultural keyword tags used at inference time.
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import List, Optional

from PIL import Image

from config.settings import settings
from config.styles import StyleSpec, get_style, list_styles
from utils.image_utils import resize_to_sdxl, center_crop_to

log = logging.getLogger(__name__)

# Caption template — uses the style's prompt_tags so the LoRA learns
# to associate them with the visual style at inference time.
CAPTION_TEMPLATE = (
    "{description}, {tags}, "
    "museum-quality heritage artwork, traditional composition, "
    "intricate detail, high resolution"
)


class DatasetPreparer:
    """Prepares per-style datasets for LoRA training."""

    def __init__(self, raw_dir: Optional[Path] = None) -> None:
        self.raw_dir = raw_dir or (settings.dataset_dir / "raw")
        self.raw_dir.mkdir(parents=True, exist_ok=True)

    def prepare_style(self, style: StyleSpec, target_size: int = 1024) -> Path:
        """Prepare dataset for one style. Returns the output directory."""
        out_dir = settings.dataset_dir / style.id
        out_dir.mkdir(parents=True, exist_ok=True)

        raw_style_dir = self.raw_dir / style.id
        if not raw_style_dir.exists():
            log.warning("No raw images found at %s — creating placeholder structure", raw_style_dir)
            raw_style_dir.mkdir(parents=True, exist_ok=True)
            self._write_readme(raw_style_dir, style)
            return out_dir

        images = sorted([p for p in raw_style_dir.iterdir()
                         if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}])
        log.info("Preparing %d images for style '%s'", len(images), style.id)

        for i, img_path in enumerate(images):
            try:
                img = Image.open(img_path).convert("RGB")
                img = resize_to_sdxl(img, target=target_size)
                img = center_crop_to(img, (target_size, target_size))

                out_img = out_dir / f"{style.id}_{i:03d}.jpg"
                img.save(out_img, quality=95)

                # Write caption
                caption = CAPTION_TEMPLATE.format(
                    description=style.description.split(".")[0].strip(),
                    tags=", ".join(style.prompt_tags),
                )
                (out_dir / f"{style.id}_{i:03d}.txt").write_text(caption, encoding="utf-8")
            except Exception as exc:
                log.warning("Failed to process %s: %s", img_path, exc)

        # Write metadata
        meta_path = out_dir / "metadata.jsonl"
        import json
        with meta_path.open("w", encoding="utf-8") as f:
            for i in range(len(images)):
                entry = {
                    "file_name": f"{style.id}_{i:03d}.jpg",
                    "prompt": CAPTION_TEMPLATE.format(
                        description=style.description.split(".")[0].strip(),
                        tags=", ".join(style.prompt_tags),
                    ),
                }
                f.write(json.dumps(entry) + "\n")

        log.info("Prepared %d samples in %s", len(images), out_dir)
        return out_dir

    def prepare_all(self) -> List[Path]:
        """Prepare datasets for all 5 heritage styles."""
        out_dirs = []
        for style in list_styles():
            out_dirs.append(self.prepare_style(style))
        return out_dirs

    @staticmethod
    def _write_readme(raw_dir: Path, style: StyleSpec) -> None:
        readme = raw_dir / "README.md"
        readme.write_text(
            f"# Raw images for {style.display_name}\n\n"
            f"Place 30-50 high-quality reference images of {style.display_name} "
            f"({style.region}) here.\n\n"
            f"## Sources\n"
            f"- Wikipedia Commons (public domain)\n"
            f"- Museum digital archives (CC-BY)\n"
            f"- Government cultural portals\n"
            f"- Books scanned under CC0\n\n"
            f"## Style description\n{style.description}\n\n"
            f"## Cultural keywords\n{', '.join(style.cultural_keywords)}\n",
            encoding="utf-8",
        )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def _cli() -> None:
    import argparse

    p = argparse.ArgumentParser(description="Prepare per-style LoRA training datasets")
    p.add_argument("--style", default=None,
                   choices=["madhubani", "warli", "pattachitra", "mughal", "tanjore"],
                   help="Prepare only this style (default: all)")
    args = p.parse_args()

    prep = DatasetPreparer()
    if args.style:
        prep.prepare_style(get_style(args.style))
    else:
        prep.prepare_all()


if __name__ == "__main__":
    _cli()