mirror of
https://github.com/bec-project/bec_widgets.git
synced 2026-10-29 08:21:27 +01:00
334 lines
14 KiB
Python
334 lines
14 KiB
Python
"""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
|
|
|
|
|
|
def test_sensor_transform_maps_frame_to_sensor_pixels():
|
|
assert jpeg_frames.sensor_transform({}, (512, 256)) is None
|
|
assert jpeg_frames.sensor_shape({}, (512, 256)) == (512, 256)
|
|
metadata = {"sensor_shape": [4096, 2048]}
|
|
assert jpeg_frames.sensor_shape(metadata, (512, 256)) == (4096, 2048)
|
|
transform = jpeg_frames.sensor_transform(metadata, (512, 256))
|
|
# image item axes: x = row, y = column
|
|
assert transform.map(512.0, 256.0) == (4096.0, 2048.0)
|
|
|
|
|
|
def test_image_draws_preview_in_sensor_pixels(qtbot, mocked_client):
|
|
view = create_widget(qtbot, Image, client=mocked_client)
|
|
metadata = {"sensor_shape": [400, 800]}
|
|
view.on_image_update_2d({"data": np.zeros((100, 200))}, metadata)
|
|
bounds = view.main_image.mapRectToView(view.main_image.boundingRect())
|
|
assert (bounds.width(), bounds.height()) == (400, 800)
|
|
view._autorange_on_next_update = False
|
|
with mock.patch.object(view, "auto_range") as auto_range:
|
|
# a new preview resolution of the same sensor keeps the zoom
|
|
view.on_image_update_2d({"data": np.zeros((50, 100))}, metadata)
|
|
auto_range.assert_not_called()
|
|
bounds = view.main_image.mapRectToView(view.main_image.boundingRect())
|
|
assert (bounds.width(), bounds.height()) == (400, 800)
|
|
# frames without a sensor shape are drawn in their own pixels again
|
|
view.on_image_update_2d({"data": np.zeros((10, 20))}, {})
|
|
bounds = view.main_image.mapRectToView(view.main_image.boundingRect())
|
|
assert (bounds.width(), bounds.height()) == (10, 20)
|
|
|
|
|
|
def _show_frame(view, image, metadata):
|
|
content, metadata = qimage_preview(
|
|
{"data": np.frombuffer(_jpeg_bytes(image), np.uint8)}, {"encoding": "jpeg", **metadata}
|
|
)
|
|
view.on_frame(content, metadata)
|
|
|
|
|
|
def test_jpeg_view_places_preview_on_sensor(qtbot, mocked_client):
|
|
view = create_widget(qtbot, JpegView, client=mocked_client)
|
|
_show_frame(view, _gradient(h=24, w=40), {"sensor_shape": [240, 80]})
|
|
bounds = view._pixmap_item.mapRectToParent(view._pixmap_item.boundingRect())
|
|
# rows along x, columns along y, scaled up to the sensor
|
|
assert (bounds.width(), bounds.height()) == (240, 80)
|
|
assert "sensor 80 x 240" in view.info_label.text()
|
|
# another preview resolution covers the same sensor area
|
|
_show_frame(view, _gradient(h=12, w=20), {"sensor_shape": [240, 80]})
|
|
bounds = view._pixmap_item.mapRectToParent(view._pixmap_item.boundingRect())
|
|
assert (bounds.width(), bounds.height()) == (240, 80)
|
|
|
|
|
|
def test_jpeg_view_crosshair_reports_sensor_pixel(qtbot, mocked_client):
|
|
view = create_widget(qtbot, JpegView, client=mocked_client)
|
|
_show_frame(view, _gradient(h=24, w=40), {"sensor_shape": [240, 400]})
|
|
view.set_crosshair(True)
|
|
assert view.crosshair_button.isChecked()
|
|
view.update_crosshair(120.7, 33.2)
|
|
assert view.cursor_label.text() == "x 120 y 33"
|
|
assert view._v_line.value() == pytest.approx(120.7)
|
|
view.update_crosshair(-1, 10)
|
|
assert view.cursor_label.text() == "outside the sensor"
|
|
view.set_crosshair(False)
|
|
assert not view._v_line.isVisible() and view.cursor_label.text() == ""
|
|
|
|
|
|
def test_jpeg_view_rois_in_sensor_pixels(qtbot, mocked_client):
|
|
view = create_widget(qtbot, JpegView, client=mocked_client)
|
|
_show_frame(view, _gradient(h=24, w=40), {"sensor_shape": [1000, 500]})
|
|
with qtbot.waitSignal(view.rois_changed) as blocker:
|
|
rect = view.add_roi("rect")
|
|
assert blocker.args[0][0]["x"] == (450, 550) and blocker.args[0][0]["y"] == (200, 300)
|
|
circle = view.add_roi("circle")
|
|
rect.setPos((10, 20))
|
|
rect.setSize((30, 40))
|
|
regions = view.roi_regions()
|
|
assert regions[0]["x"] == (10, 40) and regions[0]["y"] == (20, 60)
|
|
assert regions[1]["shape"] == "circle" and regions[1]["radius"] == 50
|
|
assert view.roi_table.rowCount() == 2
|
|
assert view.roi_table.item(0, 1).text() == "10"
|
|
# ROIs stay where they are when the preview resolution changes
|
|
_show_frame(view, _gradient(h=12, w=20), {"sensor_shape": [1000, 500]})
|
|
assert view.roi_regions()[0]["x"] == (10, 40)
|
|
view.remove_roi(circle)
|
|
assert len(view.roi_regions()) == 1
|
|
view.clear_rois()
|
|
assert view.roi_regions() == [] and view.roi_table.isHidden()
|
|
with pytest.raises(ValueError):
|
|
view.add_roi("polygon")
|
|
|
|
|
|
@pytest.mark.parametrize("shape", [(5, 7), (300, 520), (300, 520, 3)])
|
|
def test_column_first_keeps_values(shape):
|
|
image = np.random.default_rng(0).integers(0, 255, shape, dtype=np.uint8)
|
|
result = jpeg_frames.column_first(image, tile=256)
|
|
assert np.array_equal(result, image)
|
|
assert np.swapaxes(result, 0, 1).flags["C_CONTIGUOUS"]
|
|
|
|
|
|
def test_decode_jpeg_preview_lays_frame_out_column_first():
|
|
payload = np.frombuffer(_jpeg_bytes(_gradient()), np.uint8)
|
|
content, _ = decode_jpeg_preview({"data": payload}, {"encoding": "jpeg", "channels": 1})
|
|
assert content["data"].shape == (24, 40)
|
|
assert content["data"].T.flags["C_CONTIGUOUS"]
|
|
|
|
|
|
def _noisy_8bit(h=32, w=48):
|
|
return np.random.default_rng(1).integers(0, 60, (h, w), dtype=np.uint8)
|
|
|
|
|
|
def test_image_fast_path_draws_jpeg_frames_as_they_are(qtbot, mocked_client):
|
|
view = create_widget(qtbot, Image, client=mocked_client)
|
|
frame = _noisy_8bit()
|
|
with mock.patch("bec_widgets.widgets.plots.image.image_item.ImageStats.from_data") as from_data:
|
|
view.on_image_update_2d({"data": frame}, {"encoding": "jpeg"})
|
|
from_data.assert_not_called()
|
|
image = view.main_image
|
|
assert np.shares_memory(image.image, frame) # no processed copy
|
|
assert tuple(image.levels) == (0, 255) and image.autorange
|
|
# a manual range survives later frames
|
|
image.v_range = (10, 40)
|
|
view.on_image_update_2d({"data": _noisy_8bit()}, {"encoding": "jpeg"})
|
|
assert tuple(image.levels) == (10, 40)
|
|
# auto levels from the menu go back to the full 8-bit range
|
|
view.autorange = True
|
|
assert tuple(image.levels) == (0, 255)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"setup, metadata",
|
|
[
|
|
(lambda view: setattr(view, "fast_8bit_preview", False), {"encoding": "jpeg"}),
|
|
(lambda view: None, {}), # uint8 frames that are not JPEG previews
|
|
(lambda view: setattr(view.main_image, "log", True), {"encoding": "jpeg"}),
|
|
],
|
|
)
|
|
def test_image_normal_path_when_fast_path_does_not_apply(qtbot, mocked_client, setup, metadata):
|
|
view = create_widget(qtbot, Image, client=mocked_client)
|
|
setup(view)
|
|
frame = _noisy_8bit()
|
|
view.on_image_update_2d({"data": frame}, metadata)
|
|
image = view.main_image
|
|
assert tuple(image.levels) != (0, 255)
|