orisys-app/apps/studio/tools/ui_snapshot.py

#!/usr/bin/env python3
"""Render the studio UI offscreen to PNG files for visual review.

The tool never opens a camera and never starts acquisition: it builds the real
main window, optionally injects deterministic demo data, then grabs the window.

Examples
--------
    python tools/ui_snapshot.py --theme dark --out out_dark.png
    python tools/ui_snapshot.py --theme light --sensor finger \
        --panels cal,3d --demo --size 1600x900 --out out_light.png
    python tools/ui_snapshot.py --demo --scene recording --out out_recording.png
"""

from __future__ import annotations

import argparse
import math
import os
import sys
from pathlib import Path

APP_ROOT = Path(__file__).resolve().parents[1]
REPO_ROOT = APP_ROOT.parents[1]
for _path in (APP_ROOT / "src", REPO_ROOT / "shared" / "src", APP_ROOT / "tools"):
    if _path.is_dir() and str(_path) not in sys.path:
        sys.path.insert(0, str(_path))

# Silent platform must be selected before QApplication exists.
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
if sys.platform == "win32" and "QT_QPA_FONTDIR" not in os.environ:
    fonts = Path(os.environ.get("WINDIR", r"C:\Windows")) / "Fonts"
    if fonts.is_dir():
        os.environ["QT_QPA_FONTDIR"] = str(fonts)
os.chdir(APP_ROOT)


def _replay_data(window, app) -> Path:
    """Synthesize one short recording (with visible frames) and open it."""
    import numpy as np

    from orisys_app_shared.record import ExperimentRecorder, RecordConfig

    from ui_check import _fake_packet

    root = APP_ROOT / "outputs" / "_snapshot_replay"
    if root.exists():
        import shutil

        shutil.rmtree(root, ignore_errors=True)
    root.mkdir(parents=True, exist_ok=True)
    recorder = ExperimentRecorder(RecordConfig(video=True, fps=30.0, segment_seconds=30.0))
    recording = Path(
        recorder.start(
            root=root,
            session_id=1,
            metadata={
                "device": {
                    "sensor_type": "box",
                    "config_ref": str(root / "sensor.json"),
                    "config_sha256": "snapshot",
                },
                "versions": {"sdk": "snapshot"},
                "depth": {"stored": False, "replay": "recompute"},
            },
        )
    )
    # A recognisable frame: a horizontal gradient with a bright bar, so the
    # replay's "recorded image" view is reviewable instead of black.
    height, width = 108, 192
    frame = np.zeros((height, width, 3), np.uint8)
    frame[:, :, 1] = np.linspace(20, 200, width, dtype=np.uint8)[None, :]
    frame[height // 3 : 2 * height // 3, :, 2] = 180
    for index in range(150):
        recorder.on_packet(_fake_packet(index, frame=frame))
    recorder.stop()
    return recording


def _run_scene(window, app, scene: str):
    """Drive the window into a scene that would otherwise need hardware.

    ``recording`` mocks the record channel's stats through the same single seam
    the app reads them from (``_record_stats``), so nothing is written to disk
    and no device is opened: the render is of the real widgets, not of a mockup.
    ``replay`` writes one short synthetic recording and opens it in the replay
    view, which is the only way to review the progress bar without a session.
    """
    from ui_check import _fake_packet, _StuckWorker

    if scene == "replay":
        window.brand_bar.btn_playback.click()
        for _ in range(4):
            app.processEvents()
        recording = _replay_data(window, app)
        window._load_recording(recording)
        # Park the playhead where the bar and the charts show a real position.
        window.replay.set_position(window.replay.duration_s() * 0.4)
        window.arrow_panel.set_view_mode("recorded")
        for _ in range(6):
            app.processEvents()
        return

    stuck = _StuckWorker()
    window._ensure_workers = lambda: None
    window._box_worker = stuck
    window._finger_worker = stuck
    window.toolbar.btn_run.setChecked(True)
    for _ in range(4):
        app.processEvents()
    if scene != "recording":
        return

    # Real packets so the phase reaches ``running`` and the rail holds real
    # values; the recording itself is the only mocked part.
    for index in range(300):
        stuck.emit(_fake_packet(index))
    window._on_run_tick()
    for _ in range(4):
        app.processEvents()
    window._run_timer.stop()

    from orisys_app_shared.record.experiment import RecordStats

    from orisys_studio.ui import record_status

    directory = "20240101T101500Z_session0001_a1b2c3d4e5f6"
    window.toolbar.set_record_checked(True)
    window.toolbar.set_record_status(directory)
    window._record_stats = lambda: (
        record_status.RECORDING,
        RecordStats(
            duration_s=12.4,
            frames=372,
            segments=1,
            dropped_data=3,
            bytes_written=48_000_000,
        ),
        directory,
    )
    window._refresh_record_state()
    # A visible window inside the session, like the mock's 60-65 s of 0-65 s.
    first, last = window.session_store.extent()
    if last > first:
        window.force_plots.set_window(max(first, last - 6.0), last)
    for _ in range(4):
        app.processEvents()


def _demo_data(window) -> None:
    """Fill the plots, the calibration table and the arrow canvas."""
    from orisys_app_shared.force.session import SessionSample, SessionStore

    plots = window.force_plots
    # The panel reads a *selection* out of the attached source, so the demo needs
    # a real session store behind it (a store-less panel can only follow live).
    store = SessionStore()
    plots.set_store(store)
    for i in range(240):
        t = i / 30.0
        fn = 6.0 + 5.0 * math.sin(t * 1.4)
        fx = 4.0 * math.sin(t * 2.1 + 0.4)
        fy = -3.0 + 2.5 * math.cos(t * 1.1)
        store.append(
            SessionSample(
                frame_id=i,
                t_rel_s=t,
                fn_raw=fn * 7.9 + 9.0,
                fx_raw=fx * 7.9,
                fy_raw=fy * 7.9,
                fn_n=fn,
                fx_n=fx,
                fy_n=fy,
                calibrated_axes=("fn", "fx", "fy"),
            )
        )
        plots.push(t, fn, fx, fy, unit="N")

    # The overview reads the *source*, and ``push`` only rebuilds it once a second
    # (the demo fills a whole session in a few milliseconds), so build it once here.
    plots._overview_dirty = True
    plots._last_overview_at = 0.0
    plots._sync_overview()

    cal = window.cal_panel
    cal.set_running(True)
    cal.push_raw(9.81, -0.42, 0.10)
    for weight, raw, std in ((1.0, 9.72, 0.031), (2.0, 19.51, 0.044), (5.0, 48.90, 0.061)):
        from orisys_app_shared.force.calibration_io import CalSample

        cal._store.upsert("fn", CalSample(weight_n=weight, raw_mean=raw, n=180, std=std))
    cal._refresh_table()

    # Synthetic outward flow so the branded canvas and flow ramp are reviewable.
    import numpy as np

    from orisys_app_shared.acquisition.arrows import (
        draw_arrow_display,
        flow_colormap,
        hex_to_bgr,
    )
    from orisys_app_shared.theme import color

    from orisys_studio.app_window import _bgr_to_qimage

    size = 320
    grid_y, grid_x = np.mgrid[0:size, 0:size].astype(np.float32)
    centre = size / 2.0
    flow = np.zeros((size, size, 2), np.float32)
    flow[..., 0] = (grid_x - centre) * 0.10
    flow[..., 1] = (grid_y - centre) * 0.10
    blank = np.zeros((size, size, 3), np.uint8)
    rendered = draw_arrow_display(
        blank,
        flow,
        background="pure_color",
        canvas_bgr=hex_to_bgr(color(window._theme, "arrow_canvas", "#0B1830")),
        draw_style="finger" if window.sensor_picker.sensor_type() == "finger" else "box",
        colormap=flow_colormap(window._theme),
    )
    window.arrow_panel.show_frame(_bgr_to_qimage(rendered))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--theme", choices=("dark", "light"), default="dark")
    parser.add_argument("--lang", choices=("zh_CN", "en_US"), default="zh_CN")
    parser.add_argument("--sensor", choices=("box", "mini", "finger"), default="box")
    parser.add_argument(
        "--panels",
        default="",
        help="comma list of optional panels to show: cal,3d",
    )
    parser.add_argument("--size", default="1440x880", help="WxH window size")
    parser.add_argument("--view", default="arrows_pure")
    parser.add_argument(
        "--preset",
        choices=("custom", "live", "replay", "calibration"),
        default="custom",
        help="top-level view to render (the same layouts the brand rail switches)",
    )
    parser.add_argument("--demo", action="store_true", help="inject demo plot/table data")
    parser.add_argument(
        "--scene",
        choices=("idle", "connecting", "recording", "replay"),
        default="idle",
        help=(
            "connecting renders the Stop/cancel state of a hung device open; "
            "recording renders the record strip of a running recording "
            "(mocked stats: no file is written, no device is opened); "
            "replay writes one short synthetic recording and opens it in the "
            "replay view so the progress bar can be reviewed"
        ),
    )
    parser.add_argument(
        "--select",
        default="",
        help=(
            "A,B — browse the live demo session inside that window so the "
            "selection readout (cards = peak / mean · σ) can be reviewed"
        ),
    )
    parser.add_argument("--out", required=True, help="output PNG path")
    args = parser.parse_args()

    try:
        width, height = (int(part) for part in args.size.lower().split("x", 1))
    except ValueError:
        parser.error("--size must look like 1440x880")

    selection = None
    if args.select:
        try:
            start, end = (float(part) for part in args.select.split(",", 1))
        except ValueError:
            parser.error("--select must look like 2,6")
        selection = (start, end)

    panels = {part.strip().lower() for part in args.panels.split(",") if part.strip()}

    from PyQt5.QtWidgets import QApplication

    from orisys_app_shared.i18n import set_locale
    from orisys_app_shared.theme import build_qss
    from orisys_studio.app_window import StudioMainWindow

    app = QApplication(sys.argv)
    set_locale(args.lang)
    app.setStyleSheet(build_qss(args.theme))

    prefs = {
        "language": args.lang,
        "theme": args.theme,
        "sensor_type": args.sensor,
        "show_cal_panel": "cal" in panels,
        "show_finger_3d": "3d" in panels,
        "show_fps": True,
        "image_view_mode": args.view,
        "view_preset": args.preset,
    }
    window = StudioMainWindow(prefs=prefs)
    window.resize(width, height)
    window.show()
    app.processEvents()

    if args.demo and args.scene == "replay":
        pass  # the recording is the demo data

    if args.scene == "connecting":
        # A device open that never returns: the state the Stop button must
        # survive. ``ui_check._StuckWorker`` drives the same window state.
        _run_scene(window, app, "connecting")

    if args.scene == "recording":
        _run_scene(window, app, "recording")

    if args.scene == "replay":
        _run_scene(window, app, "replay")

    if args.demo and args.scene != "replay":
        _demo_data(window)
        for _ in range(4):
            app.processEvents()

    if selection is not None:
        # A selection is a browse window narrower than the session: the value
        # cards state its peak and mean · σ, the strip only its size.
        window.force_plots.set_window(*selection)
        for _ in range(4):
            app.processEvents()

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    ok = window.grab().save(str(out))
    print(f"{'saved' if ok else 'FAILED'} {out} ({width}x{height}, theme={args.theme})")
    window.close()
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())