"""AngleForge — Gradio Space + API for robotic-arm multi-angle datasets. Primary use: a (simulated) robotic arm calls the ``grab_viewpoints`` API with a real-world image and receives a series of angle/viewpoint renders to pull into a dataset. A UI is also provided to assemble, download, and publish full image datasets to Hugging Face and Edge Impulse. """ from __future__ import annotations import os import shutil import tempfile from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import gradio as gr import requests from PIL import Image from src import edge_impulse from src.backends import select_backend from src.backends.zerogpu import diagnostics as zerogpu_diagnostics, on_zerogpu from src.builder import build_dataset, generate_viewpoints from src.config import ANGLE_PRESETS, DEFAULT_ANGLES, DatasetConfig from src.hf_export import export_hf_dataset, push_to_hub ENV_HF_TOKEN = os.environ.get("HF_TOKEN", "") ENV_EI_KEY = os.environ.get("EDGE_IMPULSE_API_KEY", "") ANGLE_CHOICES = [(f"{ANGLE_PRESETS[k][0]} ({k})", k) for k in ANGLE_PRESETS] # Cache one backend instance so the heavy local pipeline loads only once. _BACKEND_CACHE: dict = {} def _get_backend(hf_token: str, image_size: int, prefer: str = "auto"): key = (prefer, bool(hf_token), image_size) if key not in _BACKEND_CACHE: _BACKEND_CACHE[key] = select_backend( hf_token=hf_token or ENV_HF_TOKEN, image_size=image_size, prefer=prefer, ) return _BACKEND_CACHE[key] # --------------------------------------------------------------------------- # # API primitive: grab a series of viewpoints from one image # --------------------------------------------------------------------------- # def _coerce_image(image: Any) -> Image.Image: """Turn whatever the API/UI passed into a PIL image. ``gr.api`` clients send images as a ``FileData`` dict (``{"path": ..., "url": ...}``) or a bare path/URL string, whereas the UI passes a PIL image directly. Normalise all of these to a PIL image. """ if isinstance(image, Image.Image): return image if isinstance(image, str): if image.startswith(("http://", "https://")): resp = requests.get(image, timeout=120) resp.raise_for_status() tmp = os.path.join(tempfile.gettempdir(), f"angleforge_in_{os.getpid()}.png") with open(tmp, "wb") as fh: fh.write(resp.content) return Image.open(tmp) return Image.open(image) if isinstance(image, dict): path = image.get("path") or image.get("name") if path and os.path.exists(path): return Image.open(path) url = image.get("url") if url: return _coerce_image(url) raise gr.Error(f"Unsupported image input: {type(image).__name__}") def grab_viewpoints( image: Image.Image, angles: Optional[List[str]] = None, seed: int = 1234, image_size: int = 512, hf_token: str = "", ) -> List[Image.Image]: """Return a series of angle/viewpoint renders for a single image. Designed to be called by a robotic-arm client via ``gradio_client``: from gradio_client import Client, handle_file client = Client("eoinedge/angleforge") views = client.predict( handle_file("part.jpg"), ["top_down", "birds_eye", "rotate_left_45"], api_name="/grab_viewpoints", ) """ if image is None: raise gr.Error("Provide an input image.") image = _coerce_image(image) angles = angles or list(DEFAULT_ANGLES) backend = _get_backend(hf_token, int(image_size)) viewpoints = generate_viewpoints( backend=backend, image=image, angles=angles, seed=int(seed), ) return [vp.image for vp in viewpoints] def _grab_for_ui(image, angles, seed, image_size, hf_token): try: backend = _get_backend(hf_token, int(image_size)) views = grab_viewpoints(image, angles, seed, image_size, hf_token) gallery = [(img, ANGLE_PRESETS.get(a, (a, ""))[0]) for img, a in zip(views, angles)] status = f"Grabbed {len(views)} viewpoint(s) using backend: {backend.source}." if backend.source == "geometric_fallback": status += ( " ⚠️ This is a geometric approximation, NOT the Qwen model — " "the Space needs a GPU (ZeroGPU) or an HF token for real angle edits." ) if on_zerogpu(): status += ( "\n\nRunning on ZeroGPU but the Qwen pipeline did not load. " "Diagnostics:\n" + zerogpu_diagnostics() ) return gallery, status except Exception as exc: # noqa: BLE001 return None, f"Error: {exc}" # --------------------------------------------------------------------------- # # Class accumulator (build up labelled source images) # --------------------------------------------------------------------------- # def add_class(label: str, files, state: Dict[str, List[str]]): state = dict(state or {}) label = (label or "").strip() if not label: return state, _class_summary(state), "Enter a class label first." paths = [f.name if hasattr(f, "name") else str(f) for f in (files or [])] if not paths: return state, _class_summary(state), "Upload at least one image for the class." state.setdefault(label, []) state[label].extend(paths) return state, _class_summary(state), f"Added {len(paths)} image(s) to class '{label}'." def clear_classes(_state): return {}, _class_summary({}), "Cleared all classes." def _class_summary(state: Dict[str, List[str]]) -> List[List[str]]: return [[label, str(len(paths))] for label, paths in (state or {}).items()] # --------------------------------------------------------------------------- # # Full dataset build # --------------------------------------------------------------------------- # def build( state: Dict[str, List[str]], dataset_name: str, angles: List[str], variations: int, plain_augs: int, image_size: int, test_ratio: float, hf_token: str, prefer_backend: str, do_push_hf: bool, hf_repo_id: str, hf_private: bool, ei_api_key: str, do_upload_ei: bool, ei_allow_duplicates: bool, progress=gr.Progress(track_tqdm=False), ): logs: List[str] = [] def log(message: str) -> str: logs.append(message) return "\n".join(logs) if not state: yield "Add at least one class first.", None, "" return work_root = Path(tempfile.mkdtemp(prefix="angleforge_")) dataset_dir = work_root / "dataset" hf_dir = work_root / "hf_dataset" try: progress(0.05, desc="Selecting backend") backend = _get_backend(hf_token, int(image_size), prefer_backend) engine = backend.source yield log(f"Using backend: {engine}"), None, "" config = DatasetConfig( out_dir=str(dataset_dir), dataset_name=dataset_name or "industrial_angles", image_size=int(image_size), angles=list(angles) or list(DEFAULT_ANGLES), variations_per_angle=int(variations), plain_augmentations_per_image=int(plain_augs), test_ratio=float(test_ratio), ) progress(0.15, desc="Generating angle images") result = build_dataset(config, backend, state, progress=lambda m: logs.append(m)) yield log(f"Generated {result.total_images} images across {len(result.label_counts)} class(es)."), None, "" progress(0.7, desc="Preparing Hugging Face imagefolder") export_hf_dataset(config, result, str(hf_dir), repo_id=hf_repo_id or "your-username/your-dataset") zip_base = work_root / f"{config.dataset_name}_dataset" zip_path = shutil.make_archive(str(zip_base), "zip", str(hf_dir)) yield log(f"Created archive: {Path(zip_path).name}"), zip_path, "" token = (hf_token or "").strip() or ENV_HF_TOKEN if do_push_hf: if not token or not hf_repo_id or "/" not in (hf_repo_id or ""): log("Skipping HF push: need a token and repo id like 'username/dataset'.") else: progress(0.85, desc="Pushing to Hugging Face") url = push_to_hub(str(hf_dir), hf_repo_id, token, private=bool(hf_private)) log(f"Pushed dataset: {url}") yield "\n".join(logs), zip_path, "" ei_key = (ei_api_key or "").strip() or ENV_EI_KEY if do_upload_ei: if not ei_key: log("Skipping Edge Impulse upload: no API key provided.") else: progress(0.92, desc="Uploading to Edge Impulse") ei_result = edge_impulse.upload_dataset( dataset_dir=str(dataset_dir), api_key=ei_key, allow_duplicates=bool(ei_allow_duplicates), progress=lambda m: logs.append(m), ) log(f"Edge Impulse: {ei_result.uploaded} uploaded, {ei_result.failed} failed.") progress(1.0, desc="Done") summary = ( f"### Done\n- Backend: **{engine}**\n- Total images: **{result.total_images}**\n" + "\n".join(f"- `{k}`: {v}" for k, v in sorted(result.label_counts.items())) ) yield "\n".join(logs), zip_path, summary except Exception as exc: # noqa: BLE001 yield log(f"ERROR: {exc}"), None, f"### Failed\n\n```\n{exc}\n```" # --------------------------------------------------------------------------- # # UI # --------------------------------------------------------------------------- # with gr.Blocks(title="AngleForge — Robotic-Arm Multi-Angle Dataset Creator") as demo: gr.Markdown( """ # 📐 AngleForge ### Robotic-Arm Multi-Angle Image Dataset Creator Turn real-world photos into multi-viewpoint image datasets for **Edge Impulse** and **Hugging Face**, using **Qwen Image Edit** (top-down/overhead, bird's-eye, worm's-eye, rotations, close-up, wide-angle). A simulated robotic arm can call the **`grab_viewpoints`** API to pull a series of angle images per object. Runs on a local GPU (free) or serverless HF Inference Providers (needs a token). With **no GPU and no token** it falls back to a **geometric** approximation so the Space always works. """ ) with gr.Tab("🤖 Grab viewpoints (API)"): with gr.Row(): with gr.Column(): vp_image = gr.Image(label="Source image", type="pil", height=280) vp_angles = gr.Dropdown( choices=ANGLE_CHOICES, value=list(DEFAULT_ANGLES), multiselect=True, label="Viewpoints / angles", ) vp_seed = gr.Slider(0, 2**31 - 1, value=1234, step=1, label="Seed") vp_size = gr.Slider(256, 1024, value=512, step=64, label="Image size (longest side)") vp_token = gr.Textbox(label="HF token (for serverless backend)", type="password", placeholder="hf_...") with gr.Row(): vp_btn = gr.Button("Grab viewpoints", variant="primary") vp_status_btn = gr.Button("Check backend status") with gr.Column(): vp_gallery = gr.Gallery(label="Viewpoints", columns=3, height=420) vp_status = gr.Textbox(label="Status", interactive=False, lines=6, max_lines=30) vp_btn.click( _grab_for_ui, inputs=[vp_image, vp_angles, vp_seed, vp_size, vp_token], outputs=[vp_gallery, vp_status], api_name="grab_viewpoints_ui", ) vp_status_btn.click( lambda: zerogpu_diagnostics(), inputs=None, outputs=[vp_status], api_name="backend_status_ui", ) with gr.Tab("🗂️ Build dataset"): state = gr.State({}) with gr.Row(): with gr.Column(): gr.Markdown("### 1. Add classes") cls_label = gr.Textbox(label="Class label", placeholder="e.g. good_part") cls_files = gr.File(label="Source images", file_count="multiple", file_types=["image"]) with gr.Row(): add_btn = gr.Button("➕ Add class") clear_btn = gr.Button("🗑️ Clear") cls_table = gr.Dataframe(headers=["label", "images"], label="Classes", interactive=False) gr.Markdown("### 2. Generation") b_dataset_name = gr.Textbox(label="Dataset name", value="industrial_angles") b_angles = gr.Dropdown(choices=ANGLE_CHOICES, value=list(DEFAULT_ANGLES), multiselect=True, label="Angles") b_variations = gr.Slider(1, 5, value=1, step=1, label="Variations per angle") b_plain = gr.Slider(0, 5, value=0, step=1, label="Extra plain augmentations per image") b_size = gr.Slider(256, 1024, value=512, step=64, label="Image size") b_test = gr.Slider(0.05, 0.5, value=0.2, step=0.05, label="Test split ratio") with gr.Column(): gr.Markdown("### 3. Backend & publishing") b_prefer = gr.Radio(["auto", "local", "serverless", "geometric"], value="auto", label="Backend preference") b_token = gr.Textbox(label="HF token", type="password", placeholder="hf_... (serverless + HF push)") do_push = gr.Checkbox(label="Push dataset to Hugging Face", value=False) b_repo = gr.Textbox(label="HF dataset repo id", placeholder="username/dataset-name") b_private = gr.Checkbox(label="Private dataset", value=False) do_ei = gr.Checkbox(label="Upload to Edge Impulse", value=False) b_ei_key = gr.Textbox(label="Edge Impulse API key", type="password", placeholder="ei_...") b_ei_dupes = gr.Checkbox(label="Allow duplicates", value=False) build_btn = gr.Button("🚀 Build dataset", variant="primary") b_summary = gr.Markdown() b_download = gr.File(label="Download dataset (zip)") b_logs = gr.Textbox(label="Logs", lines=14, max_lines=30) add_btn.click(add_class, inputs=[cls_label, cls_files, state], outputs=[state, cls_table, b_logs]) clear_btn.click(clear_classes, inputs=[state], outputs=[state, cls_table, b_logs]) build_btn.click( build, inputs=[ state, b_dataset_name, b_angles, b_variations, b_plain, b_size, b_test, b_token, b_prefer, do_push, b_repo, b_private, b_ei_key, do_ei, b_ei_dupes, ], outputs=[b_logs, b_download, b_summary], ) # Programmatic API for robot-arm clients (returns a list of images). gr.api(grab_viewpoints, api_name="grab_viewpoints") # Programmatic diagnostics for the ZeroGPU pipeline (returns a status string). gr.api(lambda: zerogpu_diagnostics(), api_name="diagnostics") if __name__ == "__main__": demo.queue().launch(show_error=True)