app.py

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