Blame for app.py

"""AnyMo: geometry-aware, setup-agnostic wearable-IMU motion understanding.

Interactive demo for CRUISEResearchGroup/AnyMo (NeurIPS 2026,
"AnyMo: Geometry-Aware Setup-Agnostic Modeling of Human Motion in the Wild").

AnyMo consumes raw wearable-IMU time series — [T, S, 6] arrays whose channels are
[acc_x, acc_y, acc_z, gyro_x, gyro_y, gyro_z] (m/s², rad/s) — attached at *any*
subset of 23 canonical body segments, and outputs a free-form motion caption,
zero-shot activity recognition against candidate labels, and IMU-text retrieval.

Try one of the bundled 5-second windows from AnyMo-Bench (head + both wrists,
60 Hz), or upload your own [T, S, 6] .npy array.
"""

import spaces  # MUST come before torch / transformers
import time
from pathlib import Path

import gradio as gr
import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt

# The AnyMo pipeline is loaded eagerly at module scope (device "cuda"): on
# ZeroGPU this lets the runtime pack the weights once at boot and stream them
# into VRAM on the first request instead of reloading every call.
print("Loading AnyMo pipeline…")
import torch

from transformers import pipeline

MODEL_ID = "CRUISEResearchGroup/AnyMo"
anymo = pipeline(
    "anymo",
    model=MODEL_ID,
    trust_remote_code=True,
    device="cuda",  # ZeroGPU: string device, never an int
    dtype=torch.bfloat16,
)
print("AnyMo pipeline ready.")

SAMPLES_DIR = Path(__file__).parent / "samples"

# Friendly body-location labels offered in the UI — all resolve through the
# processor's aliases to AnyMo's 23 canonical body segments.
LOCATION_CHOICES = [
    "head", "neck", "chest", "waist", "lower back",
    "left wrist", "right wrist", "left hand", "right hand",
    "left upper arm", "right upper arm", "left forearm", "right forearm",
    "left shoulder", "right shoulder",
    "left thigh", "right thigh", "left ankle", "right ankle",
    "left foot", "right foot",
]
# The true sensor setup of every bundled AnyMo-Bench window.
SAMPLE_LOCATIONS = ["head", "left wrist", "right wrist"]

# Bundled sample windows: real 5-second IMU streams from the AnyMo-Bench-50-US
# test split (CRUISEResearchGroup/AnyMo-Bench, CC BY-NC 4.0), each recorded
# with three sensors — Head, Left Wrist, Right Wrist — at 60 Hz.
SAMPLE_LIBRARY = {
    "walking": "walking / hiking",
    "running_in_place": "jumping / hopping / running / jogging in place",
    "sitting_reading": "sitting / reading",
    "dancing": "dancing",
    "stairs": "walking up / down stairs",
    "washing_hands": "washing / drying hands / cleaning tableware",
}
DEFAULT_CANDIDATES = "\n".join(
    [
        "A person is walking forward.",
        "A person is sitting still.",
        "A person is running.",
        "A person is dancing.",
        "A person is climbing stairs.",
        "A person is washing their hands.",
    ]
)


# ---------------------------------------------------------------- helpers
def _load_imu(path) -> np.ndarray:
    """Load a [T, S, 6] IMU array from a .npy file, validating its shape.

    Args:
        path: path to a float .npy file with shape [T, S, 6] or flattened [T, S*6].

    Returns:
        Float32 array of shape [T, S, 6].

    Raises:
        ValueError: if the array does not have a supported shape.
    """
    arr = np.load(str(path))
    arr = np.asarray(arr, dtype=np.float32)
    if arr.ndim == 2 and arr.shape[-1] % 6 == 0:
        arr = arr.reshape(arr.shape[0], -1, 6)  # e.g. [T, 18] -> [T, 3, 6]
    if arr.ndim != 3 or arr.shape[-1] != 6:
        raise ValueError(
            f"Expected an IMU array of shape [T, S, 6] (or flattened [T, S*6]); got {tuple(arr.shape)}."
        )
    return arr


def _locations_for(n_sensors: int, chosen: list[str]) -> list[str]:
    """Build a location list for n sensors from the user's dropdown choices.

    AnyMo requires every sensor to map to a distinct body segment, so extra
    sensors beyond the three dropdowns are auto-assigned unused locations.

    Args:
        n_sensors: number of sensors in the IMU array.
        chosen: the (up to 3) locations picked in the UI.

    Returns:
        List of n_sensors distinct body-location labels.
    """
    chosen = [c for c in chosen if c] or ["head"]
    extras = [c for c in LOCATION_CHOICES if c not in chosen]
    locations = (chosen + extras)[:n_sensors]
    return [l if l else "head" for l in locations]


def _plot_imu(imu: np.ndarray, locations: list[str], sampling_rate: float) -> str | None:
    """Render a per-sensor accelerometer / gyroscope preview plot.

    Args:
        imu: array of shape [T, S, 6].
        locations: one body-location label per sensor.
        sampling_rate: sampling rate of the stream in Hz.

    Returns:
        Path to the rendered PNG, or None if plotting fails.
    """
    try:
        t = np.arange(imu.shape[0]) / float(sampling_rate)
        fig, axes = plt.subplots(imu.shape[1], 2, figsize=(9, 2.0 * imu.shape[1]), squeeze=False)
        for s in range(imu.shape[1]):
            name = locations[s] if s < len(locations) else f"sensor {s + 1}"
            for c, (title, sl, unit) in enumerate(
                [("accelerometer", slice(0, 3), "m/s²"), ("gyroscope", slice(3, 6), "rad/s")]
            ):
                ax = axes[s][c]
                for ch, color in zip(range(3), ["#ef4444", "#22c55e", "#3b82f6"]):
                    ax.plot(t, imu[:, s, ch], lw=0.8, color=color)
                ax.set_title(f"{name} — {title} ({unit})", fontsize=8)
                ax.set_xlim(t[0], t[-1])
                ax.tick_params(labelsize=7)
        fig.suptitle("IMU signal preview", fontsize=10)
        fig.tight_layout(rect=(0, 0, 1, 0.97))
        out = Path("/tmp") / f"imu_preview_{int(time.time() * 1000) % 1000000}.png"
        fig.savefig(out, dpi=110)
        plt.close(fig)
        return str(out)
    except Exception as exc:  # preview is best-effort
        print("preview error:", exc)
        return None


def _parse_candidates(text: str) -> list[str]:
    """Parse the candidate-label textbox into a clean, de-duplicated list.

    Args:
        text: newline- or comma-separated candidate labels.

    Returns:
        List of non-empty candidate strings.
    """
    labels = [ln.strip() for ln in str(text).replace(",", "\n").splitlines()]
    seen, out = set(), []
    for label in labels:
        if label and label.lower() not in seen:
            seen.add(label.lower())
            out.append(label)
    return out


def _plot_scores(scores: dict[str, float], best: str) -> str:
    """Render the zero-shot recognition similarity bar chart.

    Args:
        scores: mapping from candidate label to similarity score.
        best: the winning label, highlighted in the chart.

    Returns:
        Path to the rendered PNG.
    """
    items = sorted(scores.items(), key=lambda kv: kv[1])
    labels = [k for k, _ in items]
    values = [v for _, v in items]
    fig, ax = plt.subplots(figsize=(8, max(2.0, 0.42 * len(labels))))
    ax.barh(labels, values, color=["#16a34a" if l == best else "#6366f1" for l in labels])
    ax.set_xlabel("Similarity", fontsize=9)
    ax.tick_params(labelsize=8)
    ax.set_title("Zero-shot recognition — IMU-text similarity", fontsize=10)
    fig.tight_layout()
    out = Path("/tmp") / f"anymo_scores_{int(time.time() * 1000) % 1000000}.png"
    fig.savefig(out, dpi=110)
    plt.close(fig)
    return str(out)


# ---------------------------------------------------------------- handlers
@spaces.GPU(duration=10)  # measured 3.8 s on ZeroGPU + margin for longer windows
def analyze(
    sample_selector: str = "walking",
    imu=None,
    loc0: str = "head",
    loc1: str = "left wrist",
    loc2: str = "right wrist",
    sampling_rate: float = 60.0,
    candidates_text: str = DEFAULT_CANDIDATES,
    max_new_tokens: int = 64,
):
    """Caption an IMU motion stream and classify it against candidate activities.

    Runs AnyMo's motion captioning (greedy decoding) and zero-shot activity
    recognition (IMU-text embedding similarity) over a raw wearable-IMU window.

    Args:
        sample_selector: bundled AnyMo-Bench sample to use when no upload is given.
        imu: uploaded .npy file path holding a [T, S, 6] (or [T, S*6]) IMU array;
            takes priority over the bundled sample.
        loc0: body location of sensor 1.
        loc1: body location of sensor 2.
        loc2: body location of sensor 3.
        sampling_rate: sampling rate of the IMU stream in Hz (resampled to 60).
        candidates_text: candidate activity labels, newline- or comma-separated.
        max_new_tokens: maximum number of new tokens in the generated caption.

    Returns:
        Tuple of (caption, best activity, similarity chart, preview plot, meta line).
    """
    t0 = time.time()
    if imu is None:
        imu = _load_imu(SAMPLES_DIR / f"{sample_selector}.npy")
    else:
        imu = _load_imu(imu)

    n_sensors = imu.shape[1]
    locations = _locations_for(n_sensors, [loc0, loc1, loc2])

    rate = float(sampling_rate) if sampling_rate and float(sampling_rate) > 0 else 60.0
    candidates = _parse_candidates(candidates_text) or _parse_candidates(DEFAULT_CANDIDATES)

    caption = anymo.caption(imu, locations, sampling_rate=rate, max_new_tokens=int(max_new_tokens))
    result = anymo.classify(imu, sensor_locations=locations, sampling_rate=rate, candidate_labels=candidates)
    best = result["label"]

    preview_png = _plot_imu(imu, locations, rate)
    scores_png = _plot_scores(result["scores"], best)
    elapsed = time.time() - t0
    meta = (
        f"*Sensors:* {', '.join(locations)} · *window:* {imu.shape[0]} steps @ {rate:g} Hz"
        f" · *inference:* {elapsed:.1f} s"
    )
    return str(caption), best, scores_png, preview_png, meta


# ---------------------------------------------------------------- UI
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

DESCRIPTION = """
AnyMo turns **raw wearable-IMU signals** into **motion captions** and
**zero-shot activity recognition** — with any subset of the 23 canonical body
segments, in any order. Pick a real 5-second window from AnyMo-Bench, or upload
your own `[T, S, 6]` `.npy` array (channels `[acc_x, acc_y, acc_z, gyro_x,
gyro_y, gyro_z]`, m/s² and rad/s), set each sensor's body location, and press
**Analyze motion**.

[Paper](https://arxiv.org/abs/2605.22715) · [Code](https://github.com/Breezelled/AnyMo) · [Model](https://huggingface.co/CRUISEResearchGroup/AnyMo) · [AnyMo-Bench](https://huggingface.co/datasets/CRUISEResearchGroup/AnyMo-Bench)
"""

with gr.Blocks(title="AnyMo") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"# 🏃 AnyMo — motion understanding from wearable IMUs\n{DESCRIPTION}")

        with gr.Row():
            with gr.Column(scale=1):
                sample_selector = gr.Dropdown(
                    choices=list(SAMPLE_LIBRARY.keys()),
                    value="walking",
                    label="Bundled sample (AnyMo-Bench window, 5 s @ 60 Hz)",
                    info="Real Head + Left Wrist + Right Wrist IMU windows from the AnyMo-Bench test split.",
                )
                upload = gr.File(
                    label="…or upload your own .npy ([T, S, 6] or [T, S×6])",
                    file_types=[".npy"],
                    type="filepath",
                )
                with gr.Accordion("Sensor locations", open=True):
                    loc0 = gr.Dropdown(LOCATION_CHOICES, value="head", label="Sensor 1")
                    loc1 = gr.Dropdown(LOCATION_CHOICES, value="left wrist", label="Sensor 2")
                    loc2 = gr.Dropdown(LOCATION_CHOICES, value="right wrist", label="Sensor 3")
                with gr.Accordion("Advanced", open=False):
                    sampling_rate = gr.Number(60.0, label="Sampling rate (Hz)", minimum=1.0)
                    candidates_text = gr.Textbox(
                        value=DEFAULT_CANDIDATES,
                        label="Candidate activities (zero-shot recognition)",
                        lines=6,
                        info="Newline- or comma-separated; the closest match wins.",
                    )
                    max_new_tokens = gr.Slider(16, 128, value=64, step=8, label="Max caption tokens")
                run_btn = gr.Button("Analyze motion", variant="primary")

            with gr.Column(scale=1):
                preview = gr.Image(label="IMU signal preview", type="filepath", interactive=False)
                caption_out = gr.Textbox(label="Motion caption", lines=5, buttons=["copy"])
                best_out = gr.Textbox(label="Best-matching activity", buttons=["copy"])
                scores_out = gr.Image(label="Candidate similarities", type="filepath", interactive=False)
                meta_out = gr.Markdown()

        def _preview_upload(path, rate, l0, l1, l2):
            """Refresh the preview and location defaults after a file upload.

            Args:
                path: uploaded .npy file path (or None).
                rate: sampling rate in Hz.
                l0/l1/l2: currently selected sensor locations.

            Returns:
                Gradio updates for the preview image and the location dropdowns.
            """
            try:
                if path is None:
                    return gr.update(), gr.update(), gr.update(), gr.update()
                imu = _load_imu(path)
                locations = _locations_for(imu.shape[1], [l0, l1, l2])
                plot = _plot_imu(imu, locations, rate or 60.0)
                return (
                    gr.update(value=plot),
                    gr.update(value=locations[0] if imu.shape[1] > 0 else l0),
                    gr.update(value=locations[1] if imu.shape[1] > 1 else l1),
                    gr.update(value=locations[2] if imu.shape[1] > 2 else l2),
                )
            except Exception as exc:
                print("upload preview error:", exc)
                return gr.update(), gr.update(), gr.update(), gr.update()

        def _preview_sample(name, rate):
            """Preview a bundled sample and reset locations to its true setup.

            Args:
                name: bundled sample key.
                rate: sampling rate in Hz.

            Returns:
                Gradio updates for the preview image and the location dropdowns.
            """
            try:
                imu = _load_imu(SAMPLES_DIR / f"{name}.npy")
                plot = _plot_imu(imu, SAMPLE_LOCATIONS, rate or 60.0)
                return (
                    gr.update(value=plot),
                    gr.update(value=SAMPLE_LOCATIONS[0]),
                    gr.update(value=SAMPLE_LOCATIONS[1]),
                    gr.update(value=SAMPLE_LOCATIONS[2]),
                )
            except Exception as exc:
                print("sample preview error:", exc)
                return gr.update(), gr.update(), gr.update(), gr.update()

        upload.upload(
            _preview_upload,
            inputs=[upload, sampling_rate, loc0, loc1, loc2],
            outputs=[preview, loc0, loc1, loc2],
        )
        sample_selector.change(
            _preview_sample,
            inputs=[sample_selector, sampling_rate],
            outputs=[preview, loc0, loc1, loc2],
        )

        run_btn.click(
            analyze,
            inputs=[sample_selector, upload, loc0, loc1, loc2, sampling_rate, candidates_text, max_new_tokens],
            outputs=[caption_out, best_out, scores_out, preview, meta_out],
            api_name="analyze",
        )

        gr.Examples(
            examples=[[name] for name in SAMPLE_LIBRARY],
            inputs=[sample_selector],
            outputs=[caption_out, best_out, scores_out, preview, meta_out],
            fn=analyze,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples — real AnyMo-Bench windows (Head, Left Wrist, Right Wrist @ 60 Hz)",
        )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)