src/builder.py

"""Dataset builder for AngleForge.

Two entry points:

* :func:`generate_viewpoints` — the core primitive a (simulated) robotic arm
  calls to grab a *series* of angle/viewpoint images from one real-world photo.
* :func:`build_dataset` — assembles many source images and classes into an
  Edge Impulse-ready and Hugging Face-ready image dataset.
"""

from __future__ import annotations

import csv
import json
import random
import shutil
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple

from PIL import Image

from . import images as I
from .backends import ImageEditBackend
from .config import ANGLE_PRESETS, DatasetConfig, angle_prompt

ProgressFn = Callable[[str], None]


@dataclass
class Viewpoint:
    angle_key: str
    angle_label: str
    prompt: str
    image: Image.Image


@dataclass
class BuildResult:
    out_dir: str
    backend_source: str
    total_images: int
    label_counts: Dict[str, int]
    split_counts: Dict[str, int]
    angles: List[str]
    metadata_csv: str
    summary_json: str
    failed: int = 0
    warnings: List[str] = field(default_factory=list)


def generate_viewpoints(
    backend: ImageEditBackend,
    image: Image.Image,
    angles: List[str],
    seed: int = 1234,
    num_inference_steps: int = 4,
    true_guidance_scale: float = 1.0,
    progress: Optional[ProgressFn] = None,
) -> List[Viewpoint]:
    """Grab a series of camera/gripper viewpoints from a single image.

    This is the primitive exposed to the robotic-arm API: give it one photo
    and a list of angle presets, get back the transformed views.
    """
    results: List[Viewpoint] = []
    for idx, angle_key in enumerate(angles):
        if angle_key not in ANGLE_PRESETS:
            raise ValueError(f"Unknown angle preset: {angle_key}")
        label, prompt = ANGLE_PRESETS[angle_key]
        if progress:
            progress(f"Rendering viewpoint '{label}' ({idx + 1}/{len(angles)})")
        edited = backend.edit(
            image=image,
            prompt=prompt,
            seed=seed + idx,
            num_inference_steps=num_inference_steps,
            true_guidance_scale=true_guidance_scale,
        )
        results.append(Viewpoint(angle_key, label, prompt, edited))
    return results


def _reset_dirs(out_dir: Path, labels: List[str]) -> None:
    if out_dir.exists():
        shutil.rmtree(out_dir)
    for split in ("training", "testing"):
        (out_dir / "edge_impulse_upload" / split).mkdir(parents=True, exist_ok=True)
    for split in ("train", "test"):
        for label in labels:
            (out_dir / "hf_imagefolder" / split / label).mkdir(parents=True, exist_ok=True)


def build_dataset(
    config: DatasetConfig,
    backend: ImageEditBackend,
    class_images: Dict[str, List[str]],
    progress: Optional[ProgressFn] = None,
) -> BuildResult:
    """Build a full multi-angle image dataset.

    ``class_images`` maps a class label to a list of source image paths.
    """
    def log(message: str) -> None:
        if progress:
            progress(message)
        else:
            print(message)

    rng = random.Random(config.seed)
    out_dir = Path(config.out_dir).resolve()
    labels = [I.slugify(lbl) for lbl in class_images.keys()]
    _reset_dirs(out_dir, labels)

    rows: List[Dict[str, object]] = []
    warnings: List[str] = []
    failed = 0
    label_counts: Dict[str, int] = {}
    split_counts: Dict[str, int] = {}

    def save_image(img: Image.Image, label: str, angle_key: str, source_name: str, variant: str) -> None:
        uid = I.stable_hash(f"{label}|{angle_key}|{source_name}|{variant}|{rng.random()}")
        split_dir, hf_split = ("training", "train") if rng.random() >= config.test_ratio else ("testing", "test")

        filename = f"{label}.{uid}.jpg"
        ei_path = out_dir / "edge_impulse_upload" / split_dir / filename
        I.save_jpeg(img, ei_path)

        hf_path = out_dir / "hf_imagefolder" / hf_split / label / filename
        I.save_jpeg(img, hf_path)

        rows.append(
            {
                "edge_impulse_filepath": str(ei_path.relative_to(out_dir)),
                "hf_filepath": str(hf_path.relative_to(out_dir)),
                "label": label,
                "angle": angle_key,
                "source_image": source_name,
                "variant": variant,
                "width": img.width,
                "height": img.height,
                "split": hf_split,
                "source": backend.source,
            }
        )
        label_counts[label] = label_counts.get(label, 0) + 1
        split_counts[hf_split] = split_counts.get(hf_split, 0) + 1

    total_sources = sum(len(v) for v in class_images.values())
    log(f"Backend: {backend.source}. Classes: {len(class_images)}. Source images: {total_sources}.")

    for raw_label, paths in class_images.items():
        label = I.slugify(raw_label)
        for path in paths:
            source_name = Path(path).name
            try:
                source_img = I.load_rgb(path)
            except Exception as exc:  # noqa: BLE001
                failed += 1
                warnings.append(f"Could not read {path}: {exc}")
                log(f"WARNING: could not read {path}: {exc}")
                continue

            for variation in range(max(1, config.variations_per_angle)):
                seed = config.seed + variation * 1000
                try:
                    viewpoints = generate_viewpoints(
                        backend=backend,
                        image=source_img,
                        angles=config.angles,
                        seed=seed,
                        num_inference_steps=config.num_inference_steps,
                        true_guidance_scale=config.true_guidance_scale,
                        progress=lambda m: log(f"[{label}/{source_name}] {m}"),
                    )
                except Exception as exc:  # noqa: BLE001
                    failed += 1
                    warnings.append(f"Viewpoint generation failed for {source_name}: {exc}")
                    log(f"WARNING: viewpoint generation failed for {source_name}: {exc}")
                    continue

                for vp in viewpoints:
                    variant = f"{vp.angle_key}_v{variation:02d}"
                    save_image(vp.image, label, vp.angle_key, source_name, variant)
                    for aug_idx in range(config.plain_augmentations_per_image):
                        aug = I.plain_augment(vp.image, rng)
                        save_image(aug, label, vp.angle_key, source_name, f"{variant}_aug{aug_idx:02d}")

            log(f"Processed {label}/{source_name}: {label_counts.get(label, 0)} images so far.")

    # Metadata.
    metadata_csv = out_dir / "metadata.csv"
    fieldnames = list(rows[0].keys()) if rows else [
        "edge_impulse_filepath", "hf_filepath", "label", "angle", "source_image",
        "variant", "width", "height", "split", "source",
    ]
    with metadata_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    summary = {
        "dataset": config.dataset_name,
        "backend": backend.source,
        "image_size": config.image_size,
        "angles": config.angles,
        "variations_per_angle": config.variations_per_angle,
        "plain_augmentations_per_image": config.plain_augmentations_per_image,
        "total_images": len(rows),
        "labels": label_counts,
        "splits": split_counts,
        "failed": failed,
    }
    summary_json = out_dir / "dataset_summary.json"
    summary_json.write_text(json.dumps(summary, indent=2), encoding="utf-8")

    log(f"Done. Total images: {len(rows)} (failed sources/variations: {failed}).")

    return BuildResult(
        out_dir=str(out_dir),
        backend_source=backend.source,
        total_images=len(rows),
        label_counts=label_counts,
        split_counts=split_counts,
        angles=config.angles,
        metadata_csv=str(metadata_csv),
        summary_json=str(summary_json),
        failed=failed,
        warnings=warnings,
    )