File size: 5,326 Bytes
1574b3e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """
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()
|