scripts/build_sample_wardrobe.py

"""Build a pre-processed sample wardrobe from HuggingFace datasets.

Supports multiple dataset sources:
  - fashion-1k: Codatta/Fashion-1K (flat lays, needs detection)
  - second-hand: fnauman/fashion-second-hand-front-only-rgb (individual garments)

Usage:
    cd packages/wardrobe-us
    .venv/bin/python scripts/build_sample_wardrobe.py --dataset second-hand
    .venv/bin/python scripts/build_sample_wardrobe.py --dataset fashion-1k

Requires: datasets>=2.18.0 (pip install datasets)
"""

import argparse
import io
import json
import logging
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from PIL import Image

from src.detector import detect_and_crop
from src.vision import _extract_single_garment

logging.basicConfig(level=logging.INFO, format="%(name)s | %(message)s")
logger = logging.getLogger("build_sample")

SAMPLES_DIR = Path(__file__).resolve().parent.parent / "data" / "samples"
GARMENTS_DIR = SAMPLES_DIR / "garments"
CATALOG_PATH = SAMPLES_DIR / "catalog.json"

DEFAULT_TARGET = 50
TARGET_GARMENTS = DEFAULT_TARGET

DATASETS = {
    "second-hand": {
        "hf_id": "fnauman/fashion-second-hand-front-only-rgb",
        "description": "31K individual garments on uniform background (no detection needed)",
        "needs_detection": False,
    },
    "fashion-1k": {
        "hf_id": "Codatta/Fashion-1K",
        "description": "1K flat lay outfits (multi-garment, needs detection + cropping)",
        "needs_detection": True,
    },
}

# Curated indices for Fashion-1K variety
FASHION_1K_INDICES = [
    0, 5, 12, 18, 25, 33, 41, 50, 58, 67,
    75, 83, 91, 100, 110, 120, 130, 140, 150, 160,
    170, 180, 190, 200, 220, 240, 260, 280, 300, 320,
    350, 380, 400, 430, 460, 500, 550, 600, 650, 700,
]


def save_crop(garment_id: str, crop_bytes: bytes) -> str:
    """Save a crop to the samples garments directory."""
    GARMENTS_DIR.mkdir(parents=True, exist_ok=True)
    filename = f"{garment_id}.jpg"
    path = GARMENTS_DIR / filename
    path.write_bytes(crop_bytes)
    return filename


def process_with_detection(image: Image.Image, image_idx: int, catalog: list, garment_counter: int) -> int:
    """Process a flat lay image through detection + VLM. Returns updated garment counter."""
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        image.save(tmp, format="JPEG", quality=90)
        tmp_path = tmp.name

    try:
        crops = detect_and_crop(tmp_path)
    except Exception as e:
        logger.warning("Detection failed for image %d: %s", image_idx, e)
        return garment_counter

    if not crops:
        logger.info("Image %d: no garments detected, skipping", image_idx)
        return garment_counter

    logger.info("Image %d: %d crops detected", image_idx, len(crops))

    for crop_bytes in crops:
        if garment_counter >= TARGET_GARMENTS:
            break

        garment = _extract_single_garment(crop_bytes)
        if not garment:
            logger.debug("  Crop failed VLM extraction, skipping")
            continue

        garment_counter += 1
        garment_id = f"garment_{garment_counter:03d}"
        garment["id"] = garment_id

        image_ref = save_crop(garment_id, crop_bytes)
        garment["image_ref"] = image_ref

        catalog.append(garment)
        logger.info(
            "  [%d/%d] %s: %s %s (%s)",
            garment_counter, TARGET_GARMENTS,
            garment_id, garment.get("color", "?"), garment.get("type", "?"),
            garment.get("pattern", "?"),
        )

    Path(tmp_path).unlink(missing_ok=True)
    return garment_counter


def process_individual(image: Image.Image, image_idx: int, catalog: list, garment_counter: int) -> int:
    """Process a single-garment image directly with VLM (no detection needed)."""
    buf = io.BytesIO()
    image.save(buf, format="JPEG", quality=90)
    crop_bytes = buf.getvalue()

    garment = _extract_single_garment(crop_bytes)
    if not garment:
        logger.debug("Image %d: VLM extraction failed, skipping", image_idx)
        return garment_counter

    garment_counter += 1
    garment_id = f"garment_{garment_counter:03d}"
    garment["id"] = garment_id

    image_ref = save_crop(garment_id, crop_bytes)
    garment["image_ref"] = image_ref

    catalog.append(garment)
    logger.info(
        "  [%d/%d] %s: %s %s (%s)",
        garment_counter, TARGET_GARMENTS,
        garment_id, garment.get("color", "?"), garment.get("type", "?"),
        garment.get("pattern", "?"),
    )

    return garment_counter


def build_from_fashion_1k(ds) -> list[dict]:
    """Build sample wardrobe from Fashion-1K (multi-garment flat lays)."""
    catalog: list[dict] = []
    garment_counter = 0
    max_images = 40

    for i, idx in enumerate(FASHION_1K_INDICES):
        if garment_counter >= TARGET_GARMENTS:
            break
        if idx >= len(ds):
            continue
        if i >= max_images:
            break

        sample = ds[idx]
        image = sample["image"]
        if not isinstance(image, Image.Image):
            continue
        if image.mode != "RGB":
            image = image.convert("RGB")

        logger.info("--- Processing image %d (dataset idx %d) ---", i + 1, idx)
        garment_counter = process_with_detection(image, idx, catalog, garment_counter)

    # Fill remaining with sequential if needed
    if garment_counter < TARGET_GARMENTS:
        processed = set(FASHION_1K_INDICES[:max_images])
        for idx in range(len(ds)):
            if garment_counter >= TARGET_GARMENTS:
                break
            if idx in processed:
                continue

            sample = ds[idx]
            image = sample["image"]
            if not isinstance(image, Image.Image):
                continue
            if image.mode != "RGB":
                image = image.convert("RGB")

            logger.info("--- Processing image (dataset idx %d) ---", idx)
            garment_counter = process_with_detection(image, idx, catalog, garment_counter)

    return catalog


def build_from_second_hand(ds) -> list[dict]:
    """Build sample wardrobe from second-hand dataset (individual garments)."""
    catalog: list[dict] = []
    garment_counter = 0

    # Spread indices across the dataset for variety
    step = max(1, len(ds) // (TARGET_GARMENTS * 2))
    indices = list(range(0, len(ds), step))[:TARGET_GARMENTS * 2]

    for i, idx in enumerate(indices):
        if garment_counter >= TARGET_GARMENTS:
            break

        sample = ds[idx]
        image = sample.get("image") or sample.get("img")
        if not isinstance(image, Image.Image):
            continue
        if image.mode != "RGB":
            image = image.convert("RGB")

        # Resize large images to max 512px to save VLM time
        if max(image.size) > 512:
            image.thumbnail((512, 512), Image.LANCZOS)

        logger.info("--- Processing image %d/%d (dataset idx %d) ---", i + 1, len(indices), idx)
        garment_counter = process_individual(image, idx, catalog, garment_counter)

    return catalog


def main():
    parser = argparse.ArgumentParser(description="Build sample wardrobe from HuggingFace dataset")
    parser.add_argument(
        "--dataset",
        choices=list(DATASETS.keys()),
        default="second-hand",
        help="Dataset source to use (default: second-hand)",
    )
    parser.add_argument(
        "--target",
        type=int,
        default=DEFAULT_TARGET,
        help=f"Number of garments to generate (default: {DEFAULT_TARGET})",
    )
    args = parser.parse_args()

    global TARGET_GARMENTS
    TARGET_GARMENTS = args.target

    ds_config = DATASETS[args.dataset]
    logger.info("=== Building Sample Wardrobe ===")
    logger.info("Dataset: %s (%s)", args.dataset, ds_config["description"])
    logger.info("Target: %d garments", TARGET_GARMENTS)

    try:
        from datasets import load_dataset
    except ImportError:
        logger.error("'datasets' package not installed. Run: pip install datasets")
        sys.exit(1)

    logger.info("Loading %s...", ds_config["hf_id"])
    ds = load_dataset(ds_config["hf_id"], split="train")
    logger.info("Dataset loaded: %d images", len(ds))

    SAMPLES_DIR.mkdir(parents=True, exist_ok=True)
    GARMENTS_DIR.mkdir(parents=True, exist_ok=True)

    if args.dataset == "fashion-1k":
        catalog = build_from_fashion_1k(ds)
    else:
        catalog = build_from_second_hand(ds)

    # Save catalog
    with open(CATALOG_PATH, "w", encoding="utf-8") as f:
        json.dump(catalog, f, indent=2, ensure_ascii=False)

    logger.info("=== Done ===")
    logger.info("Total garments: %d", len(catalog))
    logger.info("Catalog saved: %s", CATALOG_PATH)
    logger.info("Garment images: %s", GARMENTS_DIR)

    # Summary by type
    types: dict[str, int] = {}
    for g in catalog:
        t = g.get("type", "unknown")
        types[t] = types.get(t, 0) + 1

    logger.info("Distribution by type:")
    for t, count in sorted(types.items(), key=lambda x: -x[1]):
        logger.info("  %s: %d", t, count)


if __name__ == "__main__":
    main()