"""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, )