From 41e39c08e6eaa86d3cab98da3039b08836139104 Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Wed, 21 Oct 2026 11:49:44 +0000 Subject: [PATCH] feat(image): decode jpeg preview frames off the gui thread with a demo --- bec_widgets/examples/jpeg_preview/README.md | 29 + bec_widgets/examples/jpeg_preview/__init__.py | 0 bec_widgets/examples/jpeg_preview/demo.py | 624 ++++++++++++++++++ .../examples/jpeg_preview/jpeg_view.py | 165 +++++ .../jpeg_preview/test_offline_benchmark.py | 116 ++++ bec_widgets/utils/bec_dispatcher.py | 61 +- bec_widgets/utils/jpeg_frames.py | 132 ++++ bec_widgets/widgets/plots/image/image.py | 9 + pyproject.toml | 1 + tests/unit_tests/test_jpeg_frames.py | 191 ++++++ 10 files changed, 1324 insertions(+), 4 deletions(-) create mode 100644 bec_widgets/examples/jpeg_preview/README.md create mode 100644 bec_widgets/examples/jpeg_preview/__init__.py create mode 100644 bec_widgets/examples/jpeg_preview/demo.py create mode 100644 bec_widgets/examples/jpeg_preview/jpeg_view.py create mode 100644 bec_widgets/examples/jpeg_preview/test_offline_benchmark.py create mode 100644 bec_widgets/utils/jpeg_frames.py create mode 100644 tests/unit_tests/test_jpeg_frames.py diff --git a/bec_widgets/examples/jpeg_preview/README.md b/bec_widgets/examples/jpeg_preview/README.md new file mode 100644 index 00000000..b35ee46f --- /dev/null +++ b/bec_widgets/examples/jpeg_preview/README.md @@ -0,0 +1,29 @@ +# JPEG live preview proof of concept + +Shows detector preview frames that arrive JPEG-compressed, next to the same frames as raw counts. + +- `ophyd_devices.sim.sim_jpeg_camera.SimJpegCamera` (ophyd_devices branch `feat/jpeg-preview-poc`) + streams frames through its `preview` signal as JPEG bytes (`encoding: "jpeg"` in the message + metadata) or as raw uint16. Size, encoding, JPEG quality, rate and the 8-bit scaling change + live: `dev.jpeg_cam.configure_stream(width=4096, height=4096, encoding="jpeg_gray")`. +- The Image widget decodes JPEG previews on the thread that receives Redis messages + (`BECDispatcher.connect_slot(..., preprocess=decode_jpeg_preview)`) and keeps only the newest + frame when it falls behind (`latest_only=True`). Everything else in the widget is unchanged and + works on the decoded 0–255 values. +- `jpeg_view.JpegView` is the picture-only path: Qt decodes the JPEG into a QImage off the GUI + thread and the view draws it as a pixmap. No levels, histogram, ROI or pixel values. +- `demo.py` puts both views side by side with live camera controls, a statistics table and a + benchmark sweep (sizes x encodings x views, saved as CSV and JSON). +- `test_offline_benchmark.py` measures the same path without Redis. + +Install `simplejpeg` (`pip install -e .[jpeg]`) for the fast decoder; without it the widgets fall +back to Qt's JPEG plugin and the camera to Pillow. + +Run, with the camera loaded from `ophyd_devices/configs/jpeg_camera_example.yaml`: + +```bash +python -m bec_widgets.examples.jpeg_preview.demo --device jpeg_cam +python -m bec_widgets.examples.jpeg_preview.demo --device jpeg_cam --benchmark --out /tmp/jpeg_bench +QT_QPA_PLATFORM=offscreen python -m pytest -s -p no:randomly \ + bec_widgets/examples/jpeg_preview/test_offline_benchmark.py +``` diff --git a/bec_widgets/examples/jpeg_preview/__init__.py b/bec_widgets/examples/jpeg_preview/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bec_widgets/examples/jpeg_preview/demo.py b/bec_widgets/examples/jpeg_preview/demo.py new file mode 100644 index 00000000..3f00238b --- /dev/null +++ b/bec_widgets/examples/jpeg_preview/demo.py @@ -0,0 +1,624 @@ +""" +JPEG live preview demo: the BEC Image widget and a picture-only view side by side, fed by the +simulated JPEG camera (ophyd_devices.sim.sim_jpeg_camera.SimJpegCamera). + +Run against a BEC with the camera loaded: + + python -m bec_widgets.examples.jpeg_preview.demo --device jpeg_cam + +The top bar changes the camera stream while it runs (size, JPEG or raw, quality, rate, 8-bit +scaling). The bottom table shows, per view, what each frame costs. "Run benchmark" sweeps sizes, +encodings and views and saves the results as CSV and JSON. +""" + +from __future__ import annotations + +import argparse +import csv +import itertools +import json +import statistics +import sys +import threading +import time +from collections import deque +from dataclasses import dataclass, field +from pathlib import Path + +from bec_lib.logger import bec_logger +from qtpy.QtCore import QObject, Qt, QTimer, Signal +from qtpy.QtWidgets import ( + QApplication, + QCheckBox, + QComboBox, + QDialog, + QDoubleSpinBox, + QHBoxLayout, + QHeaderView, + QLabel, + QMainWindow, + QPushButton, + QSpinBox, + QSplitter, + QTableWidget, + QTableWidgetItem, + QVBoxLayout, + QWidget, +) + +from bec_widgets.examples.jpeg_preview.jpeg_view import JpegView +from bec_widgets.utils.colors import apply_theme +from bec_widgets.utils.error_popups import SafeSlot +from bec_widgets.widgets.plots.image.image import Image + +logger = bec_logger.logger + +RESOLUTIONS = { + "512 x 512": (512, 512), + "1024 x 1024": (1024, 1024), + "2048 x 2048": (2048, 2048), + "2068 x 2162 (Eiger 4M)": (2068, 2162), + "3108 x 3262 (Eiger 9M)": (3108, 3262), + "4096 x 4096": (4096, 4096), + "Custom": None, +} +ENCODINGS = {"JPEG gray": "jpeg_gray", "JPEG colour (colormapped)": "jpeg_rgb", "Raw uint16": "raw"} +STAT_COLUMNS = ( + "frames/s shown", + "frames/s sent", + "dropped %", + "frame KB", + "vs raw", + "encode ms (device)", + "decode ms (receiver thread)", + "GUI ms per frame", + "GUI busy %", + "latency ms", + "GUI stall ms (max)", +) + + +class TimedImage(Image): + """Image widget that records how long its update slot takes.""" + + frame_timed = Signal(dict, float) + + @SafeSlot(dict, dict) + def on_image_update_2d(self, msg: dict, metadata: dict): + t0 = time.perf_counter() + super().on_image_update_2d(msg, metadata) + # Paint now, so the measured time is the whole cost of showing this frame + self.plot_widget.viewport().repaint() + self.frame_timed.emit(metadata, t0) + + +@dataclass +class FrameSample: + """What one displayed frame cost.""" + + t_shown: float + frame_id: int + transfer_nbytes: int + raw_nbytes: int + encode_ms: float + decode_ms: float + gui_ms: float + latency_ms: float + + +@dataclass +class StatsTracker: + """Rolling statistics of one view.""" + + window_s: float = 3.0 + samples: deque = field(default_factory=lambda: deque(maxlen=2000)) + + def add(self, sample: FrameSample) -> None: + """Add one displayed frame.""" + self.samples.append(sample) + + def clear(self) -> None: + """Forget all samples.""" + self.samples.clear() + + def summary(self, since: float | None = None) -> dict: + """ + Summarise the samples of the last window (or since a time). + + Args: + since (float | None): perf_counter time to start from; default is the rolling window. + + Returns: + dict: One value per STAT_COLUMNS entry, empty if there are no samples. + """ + now = time.perf_counter() + start = since if since is not None else now - self.window_s + rows = [s for s in self.samples if s.t_shown >= start] + if len(rows) < 2: + return {} + span = max(1e-6, rows[-1].t_shown - rows[0].t_shown) + ids = [r.frame_id for r in rows] + sent = max(1, ids[-1] - ids[0]) + med = statistics.median + transfer = med(r.transfer_nbytes for r in rows) + raw = med(r.raw_nbytes for r in rows) + return { + "frames/s shown": (len(rows) - 1) / span, + "frames/s sent": sent / span, + "dropped %": max(0.0, 100.0 * (1 - (len(rows) - 1) / sent)), + "frame KB": transfer / 1e3, + "vs raw": raw / transfer if transfer else 0.0, + "encode ms (device)": med(r.encode_ms for r in rows), + "decode ms (receiver thread)": med(r.decode_ms for r in rows), + "GUI ms per frame": med(r.gui_ms for r in rows), + "latency ms": med(r.latency_ms for r in rows), + "GUI busy %": 100.0 * sum(r.gui_ms for r in rows[1:]) / 1e3 / span, + } + + +class EventLoopProbe(QObject): + """Measures how late a 20 ms timer fires: how long the GUI thread was blocked.""" + + def __init__(self, parent=None): + super().__init__(parent) + self._interval = 0.02 + self._last = time.perf_counter() + self.stalls: deque = deque(maxlen=1000) + self._timer = QTimer(self) + self._timer.timeout.connect(self._tick) + self._timer.start(int(self._interval * 1000)) + + def _tick(self): + now = time.perf_counter() + self.stalls.append((now, max(0.0, (now - self._last - self._interval) * 1e3))) + self._last = now + + def max_stall(self, since: float) -> float: + """Longest stall since a perf_counter time, in ms.""" + values = [s for t, s in self.stalls if t >= since] + return max(values) if values else 0.0 + + +class CameraControl(QWidget): + """Top bar that changes the simulated camera's stream while it runs.""" + + settings_received = Signal(dict) + rpc_failed = Signal(str) + + def __init__(self, dev, device: str, parent=None): + super().__init__(parent) + self._dev = dev + self._device = device + self._applying = False + + self.live_button = QPushButton("Stop live", self) + self.live_button.setCheckable(True) + self.live_button.setChecked(True) + self.live_button.toggled.connect(self._on_live_toggled) + self.resolution_box = QComboBox(self) + self.resolution_box.addItems(RESOLUTIONS) + self.width_spin = self._spin(16, 8192, 2048, " px") + self.height_spin = self._spin(16, 8192, 2048, " px") + self.encoding_box = QComboBox(self) + self.encoding_box.addItems(ENCODINGS) + self.quality_spin = self._spin(1, 100, 90, "") + self.rate_spin = QDoubleSpinBox(self) + self.rate_spin.setRange(0.1, 100) + self.rate_spin.setValue(5) + self.rate_spin.setSuffix(" Hz") + self.max_counts_spin = self._spin(0, 65535, 0, "") + self.max_counts_spin.setSpecialValueText("auto") + self.max_counts_spin.setToolTip("Counts mapped to 255 in JPEG modes; auto = 99.9th pct") + self.status_label = QLabel("", self) + + layout = QHBoxLayout(self) + layout.setContentsMargins(4, 4, 4, 4) + for label, widget in ( + (None, self.live_button), + ("Size", self.resolution_box), + (None, self.width_spin), + ("x", self.height_spin), + ("Encoding", self.encoding_box), + ("JPEG quality", self.quality_spin), + ("Rate", self.rate_spin), + ("255 =", self.max_counts_spin), + ): + if label: + layout.addWidget(QLabel(label, self)) + layout.addWidget(widget) + layout.addStretch(1) + layout.addWidget(self.status_label) + + self._debounce = QTimer(self) + self._debounce.setSingleShot(True) + self._debounce.setInterval(300) + self._debounce.timeout.connect(self.apply) + self.resolution_box.currentTextChanged.connect(self._on_preset) + for spin in (self.width_spin, self.height_spin, self.quality_spin, self.max_counts_spin): + spin.valueChanged.connect(self._schedule) + self.rate_spin.valueChanged.connect(self._schedule) + self.encoding_box.currentTextChanged.connect(self._schedule) + self.settings_received.connect(self._show_settings) + self.rpc_failed.connect(lambda text: self.status_label.setText(f"Error: {text}")) + self.resolution_box.setCurrentText("2048 x 2048") + + def _spin(self, low: int, high: int, value: int, suffix: str) -> QSpinBox: + spin = QSpinBox(self) + spin.setRange(low, high) + spin.setValue(value) + spin.setSuffix(suffix) + spin.setKeyboardTracking(False) + return spin + + @SafeSlot(str) + def _on_preset(self, text: str): + size = RESOLUTIONS.get(text) + custom = size is None + self.width_spin.setEnabled(custom) + self.height_spin.setEnabled(custom) + if size: + self._applying = True + self.width_spin.setValue(size[0]) + self.height_spin.setValue(size[1]) + self._applying = False + self._schedule() + + def _schedule(self, *_): + if not self._applying: + self._debounce.start() + + def current_settings(self) -> dict: + """The settings shown in the bar.""" + return { + "width": self.width_spin.value(), + "height": self.height_spin.value(), + "encoding": ENCODINGS[self.encoding_box.currentText()], + "quality": self.quality_spin.value(), + "frame_rate": self.rate_spin.value(), + "max_counts": float(self.max_counts_spin.value()), + } + + def set_settings(self, **settings) -> None: + """Change the bar (and so the camera) programmatically, e.g. from the benchmark.""" + self._applying = True + if "width" in settings or "height" in settings: + self.resolution_box.setCurrentText("Custom") + self.width_spin.setValue(settings.get("width", self.width_spin.value())) + self.height_spin.setValue(settings.get("height", self.height_spin.value())) + if "encoding" in settings: + label = {v: k for k, v in ENCODINGS.items()}[settings["encoding"]] + self.encoding_box.setCurrentText(label) + if "quality" in settings: + self.quality_spin.setValue(settings["quality"]) + if "frame_rate" in settings: + self.rate_spin.setValue(settings["frame_rate"]) + self._applying = False + self.apply() + + def _call_device(self, method: str, **kwargs): + """Call a device method off the GUI thread; the RPC waits for the device server.""" + + def run(): + try: + result = getattr(self._dev[self._device], method)(**kwargs) + if isinstance(result, dict): + self.settings_received.emit(result) + except Exception as exc: # pylint: disable=broad-except + self.rpc_failed.emit(str(exc)) + + threading.Thread(target=run, daemon=True).start() + + @SafeSlot() + def apply(self): + """Send the bar's settings to the camera.""" + self.status_label.setText("Applying…") + self._call_device("configure_stream", **self.current_settings()) + + @SafeSlot(bool) + def _on_live_toggled(self, live: bool): + self.live_button.setText("Stop live" if live else "Start live") + self._call_device("start_live" if live else "stop_live") + + @SafeSlot(dict) + def _show_settings(self, settings: dict): + self.status_label.setText( + f"camera: {settings['width']} x {settings['height']}, {settings['encoding']}" + ) + + +class JpegPreviewDemo(QMainWindow): + """Main window of the demo.""" + + def __init__(self, device: str, signal: str = "preview", parent=None): + super().__init__(parent) + self.setWindowTitle(f"JPEG live preview: {device}.{signal}") + self.device = device + + self.image = TimedImage(parent=self, popups=False) + self.image.enable_colorbar(enabled=True, style="full") + self.jpeg_view = JpegView(parent=self) + self.control = CameraControl(self.image.dev, device, parent=self) + self.probe = EventLoopProbe(self) + self.trackers = {"Image widget": StatsTracker(), "Picture view": StatsTracker()} + + self.image_check = QCheckBox("BEC Image widget (JPEG decoded to numpy)", self) + self.view_check = QCheckBox("Picture view (Qt decodes JPEG, no data)", self) + for check in (self.image_check, self.view_check): + check.setChecked(True) + self.image_check.toggled.connect(lambda on: self.enable_view("Image widget", on)) + self.view_check.toggled.connect(lambda on: self.enable_view("Picture view", on)) + + left, right = QWidget(self), QWidget(self) + for holder, check, view in ( + (left, self.image_check, self.image), + (right, self.view_check, self.jpeg_view), + ): + column = QVBoxLayout(holder) + column.setContentsMargins(0, 0, 0, 0) + column.addWidget(check) + column.addWidget(view) + splitter = QSplitter(Qt.Orientation.Horizontal, self) + splitter.addWidget(left) + splitter.addWidget(right) + splitter.setSizes([1, 1]) + + self.stats_table = QTableWidget(len(self.trackers), len(STAT_COLUMNS), self) + self.stats_table.setHorizontalHeaderLabels(STAT_COLUMNS) + self.stats_table.setVerticalHeaderLabels(list(self.trackers)) + self.stats_table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Stretch) + self.stats_table.setMaximumHeight(100) + self.benchmark_button = QPushButton("Run benchmark", self) + self.benchmark_button.clicked.connect(lambda: self.run_benchmark()) # clicked passes a bool + self.benchmark_label = QLabel("", self) + bottom = QHBoxLayout() + bottom.addWidget(self.benchmark_button) + bottom.addWidget(self.benchmark_label, 1) + + central = QWidget(self) + layout = QVBoxLayout(central) + layout.addWidget(self.control) + layout.addWidget(splitter, 1) + layout.addWidget(self.stats_table) + layout.addLayout(bottom) + self.setCentralWidget(central) + + self.image.frame_timed.connect(lambda md, t0: self._frame_done("Image widget", md, t0)) + self.jpeg_view.frame_shown.connect(self._on_view_frame) + self.image.device = device + self.image.signal = signal + self.jpeg_view.set_source(device, signal) + + self._stats_timer = QTimer(self) + self._stats_timer.timeout.connect(self._refresh_stats) + self._stats_timer.start(500) + self.benchmark: BenchmarkRunner | None = None + + def enable_view(self, name: str, enabled: bool) -> None: + """Show or hide one view; a hidden view stops receiving frames.""" + if name == "Image widget": + self.image_check.setChecked(enabled) + self.image.setVisible(enabled) + if enabled: + self.image.device = self.device + self.image.signal = "preview" + else: + self.image.signal = "" + else: + self.view_check.setChecked(enabled) + self.jpeg_view.setVisible(enabled) + if enabled: + self.jpeg_view.set_source(self.device, "preview") + else: + self.jpeg_view.bec_dispatcher.disconnect_slot(self.jpeg_view.on_frame) + self.trackers[name].clear() + + @SafeSlot() + def _on_view_frame(self): + md = self.jpeg_view.last_frame_metadata + t0 = time.perf_counter() + # Paint now, so the measured time is the whole cost of showing this frame + self.jpeg_view.view.viewport().repaint() + t0 -= self.jpeg_view.last_update_ms / 1e3 + self._frame_done("Picture view", md, t0) + + def _frame_done(self, name: str, metadata: dict, t_slot_start: float): + now = time.perf_counter() + capture = metadata.get("timestamp_capture") + self.trackers[name].add( + FrameSample( + t_shown=now, + frame_id=int(metadata.get("frame_id", 0)), + transfer_nbytes=int(metadata.get("transfer_nbytes", metadata.get("raw_nbytes", 0))), + raw_nbytes=int(metadata.get("raw_nbytes", 0)), + encode_ms=float(metadata.get("encode_ms", 0.0)), + decode_ms=float(metadata.get("decode_ms", 0.0)), + gui_ms=(now - t_slot_start) * 1e3, + latency_ms=(time.time() - capture) * 1e3 if capture else 0.0, + ) + ) + + def summaries(self, since: float | None = None) -> dict[str, dict]: + """Current statistics of both views.""" + start = since if since is not None else time.perf_counter() - 3.0 + result = {} + for name, tracker in self.trackers.items(): + summary = tracker.summary(since) + if summary: + summary["GUI stall ms (max)"] = self.probe.max_stall(start) + result[name] = summary + return result + + @SafeSlot() + def _refresh_stats(self): + for row, summary in enumerate(self.summaries().values()): + for col, key in enumerate(STAT_COLUMNS): + value = summary.get(key) + text = "" if value is None else f"{value:.1f}" + if key == "vs raw" and value: + text = f"{value:.1f}x" + self.stats_table.setItem(row, col, QTableWidgetItem(text)) + + def run_benchmark( + self, + sizes=(512, 1024, 2048, 4096), + encodings=("raw", "jpeg_gray", "jpeg_rgb"), + frame_rate: float = 10.0, + settle_s: float = 2.0, + measure_s: float = 4.0, + out_dir: Path | None = None, + on_done=None, + ): + """Sweep sizes x encodings x views; see BenchmarkRunner.""" + self.benchmark = BenchmarkRunner( + self, sizes, encodings, frame_rate, settle_s, measure_s, out_dir, on_done + ) + self.benchmark.start() + + +class BenchmarkRunner(QObject): + """ + Steps through every size, encoding and view without blocking the GUI. + + For each step only one view is enabled, the camera is reconfigured, the stream settles, + and the view's statistics over the measuring window become one result row. + """ + + def __init__( + self, demo, sizes, encodings, frame_rate, settle_s, measure_s, out_dir, on_done=None + ): + super().__init__(demo) + self.demo = demo + self.steps = list(itertools.product(sizes, encodings, ("Image widget", "Picture view"))) + self.frame_rate = frame_rate + self.settle_s = settle_s + self.measure_s = measure_s + self.out_dir = Path(out_dir) if out_dir else Path.cwd() + self.on_done = on_done + self.rows: list[dict] = [] + self._index = -1 + self._measure_start = 0.0 + self._timer = QTimer(self) + self._timer.setSingleShot(True) + self._previous_views = (demo.image_check.isChecked(), demo.view_check.isChecked()) + + def start(self): + """Run the first step.""" + self.demo.benchmark_button.setEnabled(False) + self._next_step() + + def _next_step(self): + self._index += 1 + if self._index >= len(self.steps): + self._finish() + return + size, encoding, view = self.steps[self._index] + self.demo.benchmark_label.setText( + f"Benchmark {self._index + 1}/{len(self.steps)}: {view}, {size} x {size}, {encoding}" + ) + self.demo.enable_view("Image widget", view == "Image widget") + self.demo.enable_view("Picture view", view == "Picture view") + self.demo.control.set_settings( + width=size, height=size, encoding=encoding, frame_rate=self.frame_rate + ) + self._timer.timeout.connect(self._start_measuring) + # big frames take the simulated camera a few seconds to prepare + self._timer.start(int(self.settle_s * max(1.0, size / 2048) * 1000)) + + def _start_measuring(self): + self._timer.timeout.disconnect(self._start_measuring) + self._measure_start = time.perf_counter() + self._timer.timeout.connect(self._record) + self._timer.start(int(self.measure_s * 1000)) + + def _record(self): + self._timer.timeout.disconnect(self._record) + size, encoding, view = self.steps[self._index] + summary = self.demo.summaries(since=self._measure_start).get(view, {}) + row = {"view": view, "size": f"{size}x{size}", "encoding": encoding} + row.update({key: round(summary.get(key, float("nan")), 2) for key in STAT_COLUMNS}) + self.rows.append(row) + print(f"Benchmark row: {row}", flush=True) + self._next_step() + + def _finish(self): + stamp = time.strftime("%Y%m%d-%H%M%S") + self.out_dir.mkdir(parents=True, exist_ok=True) + csv_path = self.out_dir / f"jpeg_preview_benchmark_{stamp}.csv" + with csv_path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(self.rows[0])) + writer.writeheader() + writer.writerows(self.rows) + csv_path.with_suffix(".json").write_text(json.dumps(self.rows, indent=1)) + self.demo.benchmark_label.setText(f"Benchmark saved to {csv_path}") + self.demo.benchmark_button.setEnabled(True) + self.demo.enable_view("Image widget", self._previous_views[0]) + self.demo.enable_view("Picture view", self._previous_views[1]) + print(f"Benchmark saved to {csv_path}") + if self.on_done is not None: + self.on_done(self.rows, csv_path) + else: + show_results(self.demo, self.rows) + + +def show_results(parent, rows: list[dict]) -> None: + """Show benchmark rows in a dialog.""" + dialog = QDialog(parent) + dialog.setWindowTitle("JPEG preview benchmark") + table = QTableWidget(len(rows), len(rows[0]), dialog) + table.setHorizontalHeaderLabels(list(rows[0])) + for r, row in enumerate(rows): + for c, value in enumerate(row.values()): + table.setItem(r, c, QTableWidgetItem(str(value))) + table.resizeColumnsToContents() + layout = QVBoxLayout(dialog) + layout.addWidget(table) + dialog.resize(1400, 600) + dialog.show() + + +def main(argv=None): # pragma: no cover + """Start the demo.""" + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--device", default="jpeg_cam", help="Simulated JPEG camera device") + parser.add_argument("--signal", default="preview", help="Preview signal name") + parser.add_argument("--theme", default="dark", choices=("dark", "light")) + parser.add_argument( + "--benchmark", action="store_true", help="Run the benchmark sweep at start, then quit" + ) + parser.add_argument("--sizes", default="512,1024,2048,4096", help="Benchmark frame edges") + parser.add_argument("--encodings", default="raw,jpeg_gray,jpeg_rgb") + parser.add_argument("--rate", type=float, default=10.0, help="Benchmark frame rate (Hz)") + parser.add_argument("--measure", type=float, default=4.0, help="Seconds measured per step") + parser.add_argument("--out", default=".", help="Folder for the benchmark CSV/JSON") + parser.add_argument("--screenshot", default="", help="Save a window screenshot and quit") + args = parser.parse_args(argv) + + app = QApplication.instance() or QApplication(sys.argv) + apply_theme(args.theme) + window = JpegPreviewDemo(args.device, args.signal) + window.resize(1800, 1000) + window.show() + if args.benchmark: + QTimer.singleShot( + 3000, + lambda: window.run_benchmark( + sizes=[int(s) for s in args.sizes.split(",")], + encodings=args.encodings.split(","), + frame_rate=args.rate, + measure_s=args.measure, + out_dir=Path(args.out), + on_done=lambda rows, path: app.quit(), + ), + ) + if args.screenshot: + + def grab(): + window.grab().save(args.screenshot) + app.quit() + + QTimer.singleShot(6000, grab) + sys.exit(app.exec()) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/bec_widgets/examples/jpeg_preview/jpeg_view.py b/bec_widgets/examples/jpeg_preview/jpeg_view.py new file mode 100644 index 00000000..3ca3e332 --- /dev/null +++ b/bec_widgets/examples/jpeg_preview/jpeg_view.py @@ -0,0 +1,165 @@ +""" +Picture-only live view for JPEG preview frames. + +The frame is decoded by Qt into a QImage on the thread that receives Redis messages and drawn +as a pixmap. There is no numpy data on the GUI side, so there are no levels, histogram, ROI or +pixel values; a grayscale frame can still be coloured through an 8-bit colour table. +""" + +from __future__ import annotations + +import time + +from bec_lib.endpoints import MessageEndpoints +from qtpy.QtCore import QRectF, Qt, Signal +from qtpy.QtGui import QColor, QImage, QPainter, QPixmap +from qtpy.QtWidgets import ( + QComboBox, + QGraphicsPixmapItem, + QGraphicsScene, + QGraphicsView, + QHBoxLayout, + QLabel, + QPushButton, + QVBoxLayout, + QWidget, +) + +from bec_widgets.utils.bec_widget import BECWidget +from bec_widgets.utils.colors import Colors +from bec_widgets.utils.error_popups import SafeSlot +from bec_widgets.utils.jpeg_frames import qimage_preview + +COLORMAPS = ("gray", "magma", "viridis", "plasma", "inferno") + + +class _ZoomView(QGraphicsView): + """Graphics view with wheel zoom around the cursor.""" + + def __init__(self, scene: QGraphicsScene, parent=None): + super().__init__(scene, parent) + self.setTransformationAnchor(QGraphicsView.ViewportAnchor.AnchorUnderMouse) + self.setDragMode(QGraphicsView.DragMode.ScrollHandDrag) + self.setRenderHint(QPainter.RenderHint.SmoothPixmapTransform, False) + self.setBackgroundBrush(QColor(20, 20, 20)) + self.fit_on_next_frame = True + + def wheelEvent(self, event): # pylint: disable=invalid-name + factor = 1.25 if event.angleDelta().y() > 0 else 0.8 + self.scale(factor, factor) + self.fit_on_next_frame = False + + +class JpegView(BECWidget, QWidget): + """Live view that shows preview frames as a Qt picture, without decoding to numpy.""" + + frame_shown = Signal() + + def __init__(self, parent=None, device: str | None = None, signal: str = "preview", **kwargs): + super().__init__(parent=parent, **kwargs) + self._device = None + self._signal = signal + self._color_table: list[int] | None = None + self.last_frame_metadata: dict = {} + self.last_update_ms = 0.0 + + self._scene = QGraphicsScene(self) + self._pixmap_item = QGraphicsPixmapItem() + self._pixmap_item.setTransformationMode(Qt.TransformationMode.FastTransformation) + self._scene.addItem(self._pixmap_item) + self.view = _ZoomView(self._scene, self) + + self.colormap_box = QComboBox(self) + self.colormap_box.addItems(COLORMAPS) + self.colormap_box.setToolTip("Colour table for grayscale frames (colour JPEGs ignore it)") + self.colormap_box.currentTextChanged.connect(self.set_colormap) + fit_button = QPushButton("Fit", self) + fit_button.clicked.connect(self.fit) + self.info_label = QLabel("No frame yet", self) + + bar = QHBoxLayout() + bar.addWidget(QLabel("Colormap", self)) + bar.addWidget(self.colormap_box) + bar.addWidget(fit_button) + bar.addStretch(1) + bar.addWidget(self.info_label) + layout = QVBoxLayout(self) + layout.setContentsMargins(0, 0, 0, 0) + layout.addLayout(bar) + layout.addWidget(self.view) + + if device: + self.set_source(device, signal) + + def set_source(self, device: str, signal: str = "preview") -> None: + """ + Subscribe to a device preview signal. + + Args: + device (str): Device name. + signal (str): Preview signal name. + """ + if self._device: + self.bec_dispatcher.disconnect_slot( + self.on_frame, MessageEndpoints.device_preview(self._device, self._signal) + ) + self._device, self._signal = device, signal + self.view.fit_on_next_frame = True + self.bec_dispatcher.connect_slot( + self.on_frame, + MessageEndpoints.device_preview(device, signal), + preprocess=qimage_preview, + latest_only=True, + ) + + @SafeSlot(str) + def set_colormap(self, name: str) -> None: + """ + Colour grayscale frames through a 256-entry colour table. + + Args: + name (str): Colormap name; "gray" shows the frame as it is. + """ + if name == "gray": + self._color_table = None + return + lut = Colors.get_colormap(name).getLookupTable(nPts=256, alpha=False) + self._color_table = [QColor(int(r), int(g), int(b)).rgb() for r, g, b in lut] + + @SafeSlot() + def fit(self) -> None: + """Fit the whole frame into the view.""" + self.view.fitInView(self._pixmap_item, Qt.AspectRatioMode.KeepAspectRatio) + + @SafeSlot(dict, dict) + def on_frame(self, msg: dict, metadata: dict) -> None: + """Show a frame decoded by ``qimage_preview``.""" + qimage: QImage | None = msg.get("qimage") + if qimage is None or qimage.isNull(): + return + t0 = time.perf_counter() + if self._color_table is not None and qimage.format() == QImage.Format.Format_Grayscale8: + # Same bytes, read as palette indices: colouring costs one 256-entry table + indexed = QImage( + qimage.constBits(), + qimage.width(), + qimage.height(), + qimage.bytesPerLine(), + QImage.Format.Format_Indexed8, + ) + indexed.setColorTable(self._color_table) + pixmap = QPixmap.fromImage(indexed) + else: + pixmap = QPixmap.fromImage(qimage) + new_size = pixmap.size() != self._pixmap_item.pixmap().size() + self._pixmap_item.setPixmap(pixmap) + if new_size: + self._scene.setSceneRect(QRectF(pixmap.rect())) + if new_size or self.view.fit_on_next_frame: + self.view.fit_on_next_frame = False + self.fit() + self.last_update_ms = (time.perf_counter() - t0) * 1e3 + self.last_frame_metadata = metadata + encoding = metadata.get("encoding", "raw") + self.info_label.setText(f"{qimage.width()} x {qimage.height()} {encoding}") + self.frame_shown.emit() diff --git a/bec_widgets/examples/jpeg_preview/test_offline_benchmark.py b/bec_widgets/examples/jpeg_preview/test_offline_benchmark.py new file mode 100644 index 00000000..45e0005a --- /dev/null +++ b/bec_widgets/examples/jpeg_preview/test_offline_benchmark.py @@ -0,0 +1,116 @@ +""" +Offline benchmark of the JPEG preview path, without Redis or a running BEC. + +Each frame goes through the same steps as in a live session, one view at a time: the simulated +camera makes and encodes it, BEC serialises and deserialises the message, the receiver-thread +hook decodes it, and the view updates and paints. Run it explicitly: + + QT_QPA_PLATFORM=offscreen python -m pytest -s -p no:randomly \ + bec_widgets/examples/jpeg_preview/test_offline_benchmark.py + +Environment variables: BENCH_SIZES (default "512,1024,2048,4096"), BENCH_ENCODINGS +(default "raw,jpeg_gray,jpeg_rgb"), BENCH_FRAMES (default 10), BENCH_OUT (JSON file to write). +""" + +from __future__ import annotations + +import json +import os +import statistics +import time +from pathlib import Path + +import pytest +from bec_lib import messages +from bec_lib.serialization import MsgpackSerialization + +from bec_widgets.examples.jpeg_preview.jpeg_view import JpegView +from bec_widgets.tests.fixtures import * # noqa: F401,F403 pylint: disable=wildcard-import +from bec_widgets.tests.utils import create_widget +from bec_widgets.utils.jpeg_frames import decode_jpeg_preview, qimage_preview +from bec_widgets.widgets.plots.image.image import Image + +sim_jpeg_camera = pytest.importorskip( + "ophyd_devices.sim.sim_jpeg_camera", reason="needs the ophyd_devices SimJpegCamera" +) + +SIZES = [int(s) for s in os.environ.get("BENCH_SIZES", "512,1024,2048,4096").split(",")] +ENCODINGS = os.environ.get("BENCH_ENCODINGS", "raw,jpeg_gray,jpeg_rgb").split(",") +FRAMES = int(os.environ.get("BENCH_FRAMES", "10")) + + +def _ms(t0: float) -> float: + return (time.perf_counter() - t0) * 1e3 + + +def test_offline_benchmark(qtbot, mocked_client): + camera = sim_jpeg_camera.SimJpegCamera(name="jpeg_cam") + image = create_widget(qtbot, Image, client=mocked_client, popups=False) + image.enable_colorbar(enabled=True, style="full") + view = create_widget(qtbot, JpegView, client=mocked_client) + for widget in (image, view): + widget.resize(900, 800) + widget.show() + qtbot.waitExposed(widget) + + rows = [] + for size in SIZES: + for encoding in ENCODINGS: + camera.configure_stream(width=size, height=size, encoding=encoding) + camera.make_frame() # builds the frame source for this size + samples = {key: [] for key in ("make", "encode", "wire", "deser", "dec_img")} + samples.update({"img_gui": [], "dec_pic": [], "pic_gui": []}) + for _ in range(FRAMES): + t0 = time.perf_counter() + data, metadata = camera.make_frame() + samples["make"].append(_ms(t0)) + samples["encode"].append(metadata["encode_ms"]) + msg = messages.DevicePreviewMessage( + device="jpeg_cam", signal="preview", data=data, metadata=metadata + ) + wire = MsgpackSerialization.dumps(msg) + samples["wire"].append(len(wire) / 1e3) + t0 = time.perf_counter() + received = MsgpackSerialization.loads(wire) + samples["deser"].append(_ms(t0)) + + t0 = time.perf_counter() + content, md = decode_jpeg_preview(received.content, received.metadata) + samples["dec_img"].append(_ms(t0)) + t0 = time.perf_counter() + image.on_image_update_2d(content, md) + image.plot_widget.viewport().repaint() + samples["img_gui"].append(_ms(t0)) + + t0 = time.perf_counter() + content, md = qimage_preview(received.content, received.metadata) + samples["dec_pic"].append(_ms(t0)) + t0 = time.perf_counter() + view.on_frame(content, md) + view.view.viewport().repaint() + samples["pic_gui"].append(_ms(t0)) + med = {key: round(statistics.median(values), 2) for key, values in samples.items()} + raw_kb = size * size * 2 / 1e3 + row = { + "size": f"{size}x{size}", + "encoding": encoding, + "wire KB": med["wire"], + "vs raw": round(raw_kb / med["wire"], 1), + "device make+encode ms": med["make"], + "deserialise ms": med["deser"], + "Image: decode ms (receiver)": med["dec_img"], + "Image: GUI ms": med["img_gui"], + "Picture: decode ms (receiver)": med["dec_pic"], + "Picture: GUI ms": med["pic_gui"], + } + rows.append(row) + print(row, flush=True) + + header = list(rows[0]) + lines = ["| " + " | ".join(header) + " |", "|" + " --- |" * len(header)] + lines += ["| " + " | ".join(str(row[key]) for key in header) + " |" for row in rows] + print("\n" + "\n".join(lines)) + out = os.environ.get("BENCH_OUT") + if out: + Path(out).write_text(json.dumps(rows, indent=1)) + camera.destroy() diff --git a/bec_widgets/utils/bec_dispatcher.py b/bec_widgets/utils/bec_dispatcher.py index 02024c7c..afd396fc 100644 --- a/bec_widgets/utils/bec_dispatcher.py +++ b/bec_widgets/utils/bec_dispatcher.py @@ -3,6 +3,7 @@ from __future__ import annotations import collections import random import string +import threading import time import weakref from collections.abc import Callable @@ -79,7 +80,14 @@ class QtThreadSafeCallback(QObject): cb_signal = pyqtSignal(dict, dict) - def __init__(self, cb: Callable, cb_info: dict | None = None, owner: object | None = None): + def __init__( + self, + cb: Callable, + cb_info: dict | None = None, + owner: object | None = None, + preprocess: Callable[[dict, dict], tuple[dict, dict]] | None = None, + latest_only: bool = False, + ): """ Initialize the QtThreadSafeCallback. @@ -89,9 +97,21 @@ class QtThreadSafeCallback(QObject): Defaults to None. owner (object, optional): Lifetime anchor for non-method callables. Ignored for bound methods (their ``__self__`` is the owner). Defaults to None. + preprocess (Callable, optional): Called with ``(msg_content, metadata)`` on the + receiving thread before the message is handed to the GUI thread; returns the + new ``(msg_content, metadata)``. Use it for slow, widget-free work such as + decoding compressed frames. Defaults to None. + latest_only (bool, optional): While a message waits for the GUI thread, a newer + one replaces it instead of queueing behind it, so a slow slot shows the newest + frame instead of falling further behind. Defaults to False. """ super().__init__() self.cb_info = cb_info + self.preprocess = preprocess + self.latest_only = latest_only + self._pending: tuple[dict, dict] | None = None + self._pending_lock = threading.Lock() + self.dropped = 0 if hasattr(cb, "__self__"): # Bound method: the owner is the method's instance; hold NO strong reference @@ -118,12 +138,17 @@ class QtThreadSafeCallback(QObject): # make 2 differents QtThreadSafeCallback to look # identical when used as dictionary keys, if the # callback is the same - return f"{id(self.cb_ref)}{self.cb_info}".__hash__() + return f"{id(self.cb_ref)}{self.cb_info}{id(self.preprocess)}{self.latest_only}".__hash__() def __eq__(self, other): if not isinstance(other, QtThreadSafeCallback): return False - return self.cb_ref == other.cb_ref and self.cb_info == other.cb_info + return ( + self.cb_ref == other.cb_ref + and self.cb_info == other.cb_info + and self.preprocess is other.preprocess + and self.latest_only == other.latest_only + ) def invalidate(self) -> None: """Suppress pending deliveries and release any strongly held callback.""" @@ -135,6 +160,12 @@ class QtThreadSafeCallback(QObject): """Resolve and retain a live callback for the duration of this delivery.""" if not self._active: return + if self.latest_only: + with self._pending_lock: + pending, self._pending = self._pending, None + if pending is None: + return + msg_content, metadata = pending owner = self.cb_owner() if self.cb_owner is not None else None if self.cb_owner is not None and owner is None: return @@ -155,6 +186,19 @@ class QtThreadSafeCallback(QObject): def __call__(self, msg_content, metadata): if not self._active: return + if self.preprocess is not None: + try: + msg_content, metadata = self.preprocess(msg_content, metadata or {}) + except Exception: # pylint: disable=broad-except + logger.exception(f"Preprocessing a message for {self.cb_ref()} failed") + return + if self.latest_only: + with self._pending_lock: + replaced = self._pending is not None + self._pending = (msg_content, metadata) + if replaced: + self.dropped += 1 + return self.cb_signal.emit(msg_content, metadata) @@ -254,6 +298,8 @@ class BECDispatcher: topics: EndpointInfo | str | list[EndpointInfo] | list[str], cb_info: dict | None = None, owner: object | None = None, + preprocess: Callable[[dict, dict], tuple[dict, dict]] | None = None, + latest_only: bool = False, **kwargs, ) -> None: """Connect widget's qt slot, so that it is called on new pub/sub topic message. @@ -267,6 +313,11 @@ class BECDispatcher: owner (object | None): Lifetime anchor for non-method callables (lambdas, partials, module functions): the subscription is released when the owner is cleaned up or destroyed. Bound methods already carry their owner and ignore this. Defaults to None. + preprocess (Callable | None): Runs on the receiving thread with + ``(msg_content, metadata)`` and returns the pair the slot gets, e.g. to decode + compressed frames without blocking the GUI. Defaults to None. + latest_only (bool): Deliver only the newest message when the slot falls behind; + older undelivered ones are dropped. Meant for live previews. Defaults to False. """ if not hasattr(slot, "__self__") and owner is None: logger.warning( @@ -277,7 +328,9 @@ class BECDispatcher: # Self-healing: reap subscriptions whose owners died without any close event # before adding new ones, so stale registrations never accumulate. self.cleanup_dead_slots() - qt_slot = QtThreadSafeCallback(cb=slot, cb_info=cb_info, owner=owner) + qt_slot = QtThreadSafeCallback( + cb=slot, cb_info=cb_info, owner=owner, preprocess=preprocess, latest_only=latest_only + ) if not self.client.connector.any_stream_is_registered(topics, qt_slot): if qt_slot not in self._registered_slots: self._registered_slots[qt_slot] = qt_slot diff --git a/bec_widgets/utils/jpeg_frames.py b/bec_widgets/utils/jpeg_frames.py new file mode 100644 index 00000000..11fa9d60 --- /dev/null +++ b/bec_widgets/utils/jpeg_frames.py @@ -0,0 +1,132 @@ +""" +Helpers for preview frames that arrive JPEG-compressed. + +A JPEG preview frame is a ``DevicePreviewMessage`` whose ``data`` is a flat uint8 array holding +the JPEG bytes, with message metadata ``{"encoding": "jpeg", "shape": [h, w], "channels": 1 | 3}``. +Frames without that metadata are passed through untouched. + +The ``*_preview`` functions match the ``preprocess`` hook of ``BECDispatcher.connect_slot``: +they run on the thread that receives Redis messages, so decoding never blocks the GUI thread. +""" + +from __future__ import annotations + +import time + +import numpy as np +from qtpy.QtCore import QByteArray +from qtpy.QtGui import QImage + +try: # libjpeg-turbo bindings; much faster than Qt's own JPEG plugin + import simplejpeg +except ImportError: # pragma: no cover - depends on the environment + simplejpeg = None + + +def is_jpeg_frame(metadata: dict | None) -> bool: + """Whether a preview message carries a JPEG-compressed frame.""" + return bool(metadata) and metadata.get("encoding") == "jpeg" + + +def _as_bytes(data: bytes | np.ndarray) -> bytes: + if isinstance(data, np.ndarray): + return data.tobytes() + return bytes(data) + + +def qimage_to_array(image: QImage) -> np.ndarray: + """ + Copy an 8-bit grayscale or RGB QImage into a numpy array. + + Args: + image (QImage): The image; other formats are converted to RGB888 first. + + Returns: + np.ndarray: (h, w) uint8 for grayscale, (h, w, 3) uint8 otherwise. + """ + if image.format() not in (QImage.Format.Format_Grayscale8, QImage.Format.Format_RGB888): + image = image.convertToFormat(QImage.Format.Format_RGB888) + channels = 1 if image.format() == QImage.Format.Format_Grayscale8 else 3 + h, w, bpl = image.height(), image.width(), image.bytesPerLine() + buffer = np.frombuffer(image.constBits(), dtype=np.uint8, count=h * bpl).reshape(h, bpl) + array = buffer[:, : w * channels].reshape(h, w, channels) if channels == 3 else buffer[:, :w] + return array.copy() + + +def decode_jpeg(data: bytes | np.ndarray, channels: int | None = None) -> np.ndarray: + """ + Decode JPEG bytes into a numpy array. + + Args: + data (bytes | np.ndarray): The JPEG bytes, or a uint8 array holding them. + channels (int | None): 1 for grayscale, 3 for RGB, None to keep what the file has. + + Returns: + np.ndarray: (h, w) uint8 for grayscale, (h, w, 3) uint8 for RGB. + """ + raw = _as_bytes(data) + if simplejpeg is not None: + if channels is None: + channels = 1 if simplejpeg.decode_jpeg_header(raw)[2] == "Gray" else 3 + colorspace = "GRAY" if channels == 1 else "RGB" + image = simplejpeg.decode_jpeg(raw, colorspace=colorspace) + return image[..., 0] if channels == 1 else image + qimage = QImage.fromData(QByteArray(raw), "JPG") + if qimage.isNull(): + raise ValueError("Could not decode JPEG frame") + if channels == 1: + qimage = qimage.convertToFormat(QImage.Format.Format_Grayscale8) + return qimage_to_array(qimage) + + +def decode_jpeg_preview(msg_content: dict, metadata: dict) -> tuple[dict, dict]: + """ + Dispatcher preprocess hook: replace JPEG bytes in a preview message by the decoded array. + + Args: + msg_content (dict): Content of the DevicePreviewMessage. + metadata (dict): Metadata of the message. + + Returns: + tuple[dict, dict]: New content and metadata; the inputs are not modified. + """ + if not is_jpeg_frame(metadata): + return msg_content, metadata + data = msg_content.get("data") + t0 = time.perf_counter() + image = decode_jpeg(data, metadata.get("channels")) + decode_ms = (time.perf_counter() - t0) * 1e3 + metadata = {**metadata, "decode_ms": decode_ms, "transfer_nbytes": int(np.size(data))} + return {**msg_content, "data": image}, metadata + + +def qimage_preview(msg_content: dict, metadata: dict) -> tuple[dict, dict]: + """ + Dispatcher preprocess hook for picture-only views: decode the frame into a QImage with Qt. + + JPEG frames are decoded by Qt's JPEG plugin; raw frames are scaled to 8 bit (0 to the + 99.9th percentile). The QImage is put under ``msg_content["qimage"]``. + + Args: + msg_content (dict): Content of the DevicePreviewMessage. + metadata (dict): Metadata of the message. + + Returns: + tuple[dict, dict]: New content and metadata; the inputs are not modified. + """ + data = msg_content.get("data") + t0 = time.perf_counter() + if is_jpeg_frame(metadata): + qimage = QImage.fromData(QByteArray(_as_bytes(data)), "JPG") + else: + frame = np.asarray(data) + if frame.ndim != 2: + return msg_content, metadata + top = float(np.percentile(frame[::8, ::8], 99.9)) or 1.0 + frame8 = np.ascontiguousarray(np.clip(frame * (255.0 / top), 0, 255).astype(np.uint8)) + h, w = frame8.shape + qimage = QImage(frame8.data, w, h, frame8.strides[0], QImage.Format.Format_Grayscale8) + qimage = qimage.copy() # detach from the numpy buffer + decode_ms = (time.perf_counter() - t0) * 1e3 + metadata = {**metadata, "decode_ms": decode_ms, "transfer_nbytes": int(np.asarray(data).nbytes)} + return {**msg_content, "qimage": qimage}, metadata diff --git a/bec_widgets/widgets/plots/image/image.py b/bec_widgets/widgets/plots/image/image.py index bc606dd3..6411a153 100644 --- a/bec_widgets/widgets/plots/image/image.py +++ b/bec_widgets/widgets/plots/image/image.py @@ -12,6 +12,7 @@ from qtpy.QtWidgets import QWidget from bec_widgets.utils.bec_connector import ConnectionConfig from bec_widgets.utils.colors import Colors, apply_theme from bec_widgets.utils.error_popups import SafeProperty, SafeSlot +from bec_widgets.utils.jpeg_frames import decode_jpeg_preview from bec_widgets.widgets.plots.image.image_base import ImageBase from bec_widgets.widgets.plots.image.image_item import ImageItem from bec_widgets.widgets.plots.image.toolbar_components.device_selection import ( @@ -118,6 +119,7 @@ class Image(ImageBase): config = ImageConfig(widget_class=self.__class__.__name__) self.gui_id = config.gui_id self.subscriptions: defaultdict[str, ImageLayerConfig] = defaultdict(ImageLayerConfig) + self.last_frame_metadata: dict = {} # Store signal configs separately (not serialized to QSettings) self._signal_configs: dict[str, dict] = {} @@ -447,6 +449,8 @@ class Image(ImageBase): self.bec_dispatcher.connect_slot( self.on_image_update_2d, MessageEndpoints.device_preview(self._config.device, self._config.signal), + preprocess=decode_jpeg_preview, + latest_only=True, ) elif signal_class in self.SUPPORTED_SIGNALS: self.async_update = True @@ -1003,6 +1007,11 @@ class Image(ImageBase): if data is None: logger.warning("No data received for image update from 2D.") return + self.last_frame_metadata = metadata + previous = image.raw_data + if previous is not None and np.shape(previous) != np.shape(data): + # e.g. the detector changed resolution: show the whole new frame + self._autorange_on_next_update = True if self._color_bar is not None: self._color_bar.blockSignals(True) image.set_data(data) diff --git a/pyproject.toml b/pyproject.toml index e40ff99e..766a33a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ dev = [ "pre_commit~=4.2", ] qtermwidget = ["pyside6_qtermwidget==0.6.6.11.1"] +jpeg = ["simplejpeg~=1.8"] [build-system] diff --git a/tests/unit_tests/test_jpeg_frames.py b/tests/unit_tests/test_jpeg_frames.py new file mode 100644 index 00000000..46b18d8d --- /dev/null +++ b/tests/unit_tests/test_jpeg_frames.py @@ -0,0 +1,191 @@ +"""JPEG preview frames: decoding helpers, dispatcher hooks, Image widget and picture view.""" + +import threading +from unittest import mock + +import numpy as np +import pytest +from qtpy.QtCore import QBuffer, QByteArray, QIODevice +from qtpy.QtGui import QImage + +from bec_widgets.examples.jpeg_preview.jpeg_view import JpegView +from bec_widgets.tests.utils import create_widget +from bec_widgets.utils import jpeg_frames +from bec_widgets.utils.bec_dispatcher import QtThreadSafeCallback +from bec_widgets.utils.jpeg_frames import decode_jpeg, decode_jpeg_preview, qimage_preview +from bec_widgets.widgets.plots.image.image import Image + + +def _jpeg_bytes(image: np.ndarray, quality: int = 95) -> bytes: + """Encode with Qt so the tests do not need simplejpeg.""" + image = np.ascontiguousarray(image) + h, w = image.shape[:2] + if image.ndim == 2: + qimage = QImage(image.data, w, h, image.strides[0], QImage.Format.Format_Grayscale8) + else: + qimage = QImage(image.data, w, h, image.strides[0], QImage.Format.Format_RGB888) + data = QByteArray() + buffer = QBuffer(data) + buffer.open(QIODevice.OpenModeFlag.WriteOnly) + qimage.save(buffer, "JPG", quality) + return bytes(data) + + +def _gradient(h=24, w=40): + return np.tile(np.linspace(0, 255, w, dtype=np.uint8), (h, 1)) + + +@pytest.fixture(params=["simplejpeg", "qt"]) +def decoder(request, monkeypatch): + if request.param == "qt": + monkeypatch.setattr(jpeg_frames, "simplejpeg", None) + elif jpeg_frames.simplejpeg is None: + pytest.skip("simplejpeg not installed") + return request.param + + +def test_decode_jpeg_gray_and_rgb(decoder): + gray = _gradient() + decoded = decode_jpeg(np.frombuffer(_jpeg_bytes(gray), np.uint8), channels=1) + assert decoded.shape == gray.shape and decoded.dtype == np.uint8 + assert np.abs(decoded.astype(int) - gray).max() < 12 + + rgb = np.stack([gray, gray[:, ::-1], np.full_like(gray, 128)], axis=-1) + decoded = decode_jpeg(_jpeg_bytes(rgb), channels=3) + assert decoded.shape == rgb.shape + assert np.abs(decoded.astype(int) - rgb).mean() < 6 + + +def test_decode_jpeg_preview_replaces_data_and_keeps_inputs(decoder): + payload = np.frombuffer(_jpeg_bytes(_gradient()), np.uint8) + content = {"data": payload, "device": "cam"} + metadata = {"encoding": "jpeg", "channels": 1, "shape": [24, 40]} + new_content, new_metadata = decode_jpeg_preview(content, metadata) + assert new_content["data"].shape == (24, 40) + assert new_metadata["transfer_nbytes"] == payload.size + assert new_metadata["decode_ms"] >= 0 + assert content["data"] is payload and "decode_ms" not in metadata + + +def test_decode_jpeg_preview_passes_raw_frames_through(): + content, metadata = {"data": np.zeros((4, 4))}, {"encoding": "raw"} + assert decode_jpeg_preview(content, metadata) == (content, metadata) + + +def test_qimage_preview_jpeg_and_raw(): + content = {"data": np.frombuffer(_jpeg_bytes(_gradient()), np.uint8)} + new_content, _ = qimage_preview(content, {"encoding": "jpeg"}) + assert new_content["qimage"].size().width() == 40 + + raw = {"data": (np.arange(64, dtype=np.uint16) * 100).reshape(8, 8)} + new_content, _ = qimage_preview(raw, {}) + assert new_content["qimage"].format() == QImage.Format.Format_Grayscale8 + + +def test_callback_preprocess_runs_before_emit(qtbot): + received = [] + wrapper = QtThreadSafeCallback( + lambda content, metadata: received.append((content, metadata)), + preprocess=lambda content, metadata: ({**content, "x": 2}, {**metadata, "y": 3}), + ) + wrapper({"x": 1}, {}) + qtbot.waitUntil(lambda: bool(received)) + assert received == [({"x": 2}, {"y": 3})] + + +def test_callback_preprocess_error_drops_message(qtbot): + received = [] + + def fail(content, metadata): + raise ValueError("broken frame") + + wrapper = QtThreadSafeCallback(lambda *args: received.append(args), preprocess=fail) + wrapper({"x": 1}, {}) + qtbot.wait(20) + assert received == [] + + +def test_callback_latest_only_delivers_newest(qtbot): + received = [] + wrapper = QtThreadSafeCallback( + lambda content, metadata: received.append(content["n"]), latest_only=True + ) + # Messages arrive on the Redis thread; the GUI thread is busy until the event loop runs + sender = threading.Thread(target=lambda: [wrapper({"n": n}, {}) for n in range(5)]) + sender.start() + sender.join() + qtbot.waitUntil(lambda: bool(received)) + qtbot.wait(20) + assert received == [4] + assert wrapper.dropped == 4 + sender = threading.Thread(target=lambda: wrapper({"n": 5}, {})) + sender.start() + sender.join() + qtbot.waitUntil(lambda: received[-1] == 5) + + +def test_callbacks_with_different_hooks_are_different_wrappers(): + def slot(content, metadata): + pass + + plain = QtThreadSafeCallback(slot) + assert plain == QtThreadSafeCallback(slot) + assert plain != QtThreadSafeCallback(slot, preprocess=decode_jpeg_preview) + assert plain != QtThreadSafeCallback(slot, latest_only=True) + + +def _set_preview_signal(client, device="eiger", signal="preview"): + client.device_manager.devices[device]._info["signals"][signal] = { + "obj_name": f"{device}_{signal}", + "signal_class": "PreviewSignal", + "component_name": signal, + "describe": {"signal_info": {"ndim": 2}}, + } + + +def test_image_preview_connects_with_jpeg_decoding(qtbot, mocked_client): + view = create_widget(qtbot, Image, client=mocked_client) + _set_preview_signal(mocked_client) + with mock.patch.object(view.bec_dispatcher, "connect_slot") as connect_slot: + view.image(device="eiger", signal="preview") + kwargs = connect_slot.call_args.kwargs + assert kwargs["preprocess"] is decode_jpeg_preview + assert kwargs["latest_only"] is True + + +def test_image_shows_decoded_jpeg_frame(qtbot, mocked_client, decoder): + view = create_widget(qtbot, Image, client=mocked_client) + payload = np.frombuffer(_jpeg_bytes(_gradient()), np.uint8) + content, metadata = decode_jpeg_preview( + {"data": payload}, {"encoding": "jpeg", "channels": 1, "frame_id": 7} + ) + view.on_image_update_2d(content, metadata) + assert view.main_image.image.shape == (24, 40) + assert view.last_frame_metadata["frame_id"] == 7 + + +def test_image_autoranges_when_frame_size_changes(qtbot, mocked_client): + view = create_widget(qtbot, Image, client=mocked_client) + view.on_image_update_2d({"data": np.zeros((8, 8))}, {}) + view._autorange_on_next_update = False + with mock.patch.object(view, "auto_range") as auto_range: + view.on_image_update_2d({"data": np.ones((8, 8))}, {}) + auto_range.assert_not_called() + view.on_image_update_2d({"data": np.ones((16, 32))}, {}) + auto_range.assert_called_once() + + +def test_jpeg_view_shows_frame_and_colormap(qtbot, mocked_client): + view = create_widget(qtbot, JpegView, client=mocked_client) + content, metadata = qimage_preview( + {"data": np.frombuffer(_jpeg_bytes(_gradient()), np.uint8)}, {"encoding": "jpeg"} + ) + with qtbot.waitSignal(view.frame_shown): + view.on_frame(content, metadata) + assert view._pixmap_item.pixmap().width() == 40 + view.set_colormap("magma") + assert len(view._color_table) == 256 + view.on_frame(content, metadata) + assert view._pixmap_item.pixmap().height() == 24 + view.set_colormap("gray") + assert view._color_table is None