163 lines
6.4 KiB
Python
163 lines
6.4 KiB
Python
"""B-scan rendering and cache helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
|
|
from PyQt6.QtCore import QRectF
|
|
import numpy as np
|
|
import pyqtgraph as pg
|
|
|
|
from python_app.models.collection_identity import collection_key
|
|
from python_app.models.dataset_model import ResultCollection
|
|
from python_app.models.gui_profile_schema import GuiBscanStateModel
|
|
from python_app.gui.controllers.renderers.common import build_lut, create_plot
|
|
|
|
|
|
def _result_tail(
|
|
*,
|
|
result_history: list[ResultCollection],
|
|
history_limit: int,
|
|
) -> list[ResultCollection]:
|
|
"""Choose by acquisition order, keeping the newest result of each exact sweep."""
|
|
unique_reversed_tail: list[ResultCollection] = []
|
|
seen_keys: set[tuple[int, int]] = set()
|
|
for collection in reversed(result_history[-history_limit:]):
|
|
key = collection_key(collection)
|
|
if key in seen_keys:
|
|
continue
|
|
seen_keys.add(key)
|
|
unique_reversed_tail.append(collection)
|
|
|
|
unique_reversed_tail.reverse()
|
|
return unique_reversed_tail
|
|
|
|
|
|
def apply_mean_ascan_subtraction(
|
|
history: deque[np.ndarray],
|
|
*,
|
|
enabled: bool,
|
|
) -> np.ndarray:
|
|
"""Subtract mean A-scan across sweeps when enabled."""
|
|
sweeps = np.vstack(history).astype(np.float32, copy=False)
|
|
if not enabled:
|
|
return sweeps
|
|
mean_trace = np.mean(sweeps, axis=0, dtype=np.float32)
|
|
return sweeps - mean_trace
|
|
|
|
|
|
def rebuild_bscan_history_from_results(
|
|
result_history: list[ResultCollection],
|
|
history_limit: int,
|
|
) -> tuple[dict[tuple[int, int], deque[np.ndarray]], dict[tuple[int, int], np.ndarray]]:
|
|
"""Group compatible range profiles; a changed depth axis starts a new image."""
|
|
history_by_combo: dict[tuple[int, int], deque[np.ndarray]] = {}
|
|
depth_axis_by_combo: dict[tuple[int, int], np.ndarray] = {}
|
|
|
|
result_tail = _result_tail(
|
|
result_history=result_history,
|
|
history_limit=history_limit,
|
|
)
|
|
|
|
for collection in result_tail:
|
|
for block in collection.blocks:
|
|
key = (block.combo.input, block.combo.output)
|
|
for payload in block.payloads:
|
|
if payload.kind != 1 or payload.processing_name != "bscan":
|
|
continue
|
|
if payload.frequency_hz.size == 0 or payload.trace.size == 0:
|
|
continue
|
|
if payload.frequency_hz.size != payload.trace.size:
|
|
continue
|
|
|
|
depth_axis = np.asarray(payload.frequency_hz, dtype=np.float32)
|
|
amplitudes = np.asarray(np.real(payload.trace), dtype=np.float32)
|
|
history = history_by_combo.get(key)
|
|
stored_axis = depth_axis_by_combo.get(key)
|
|
if (
|
|
history is None
|
|
or stored_axis is None
|
|
or stored_axis.shape != depth_axis.shape
|
|
or not np.allclose(stored_axis, depth_axis, rtol=1e-4, atol=1e-6)
|
|
):
|
|
history = deque(maxlen=history_limit)
|
|
history_by_combo[key] = history
|
|
depth_axis_by_combo[key] = depth_axis.copy()
|
|
|
|
history.append(amplitudes.copy())
|
|
|
|
return history_by_combo, depth_axis_by_combo
|
|
|
|
|
|
def bscan_lookup_table(axis_mode: str) -> np.ndarray:
|
|
"""Build B-scan colormap table for selected axis mode."""
|
|
if axis_mode == "abs":
|
|
return build_lut(["#440154", "#31688e", "#35b779", "#fde725"])
|
|
return build_lut(["#2166ac", "#67a9cf", "#f7f7f7", "#ef8a62", "#b2182b"])
|
|
|
|
|
|
def bscan_levels(sweeps: np.ndarray, axis_mode: str) -> tuple[float, float]:
|
|
"""Compute image levels for B-scan data based on axis mode."""
|
|
min_value = float(np.min(sweeps))
|
|
max_value = float(np.max(sweeps))
|
|
if axis_mode == "abs":
|
|
if max_value <= min_value:
|
|
return min_value, min_value + 1e-6
|
|
return min_value, max_value
|
|
|
|
max_abs = max(abs(min_value), abs(max_value), 1e-6)
|
|
return -max_abs, max_abs
|
|
|
|
|
|
class BscanRenderer:
|
|
"""Own the heatmap and its per-combination arrays; input history stays external."""
|
|
|
|
def __init__(self) -> None:
|
|
self.widget = create_plot()
|
|
self.widget.setLabel("bottom", "Sweep #")
|
|
self.widget.setLabel("left", "Range", units="m")
|
|
item = self.widget.getPlotItem()
|
|
assert item is not None
|
|
item.getViewBox().invertY(False)
|
|
item.getViewBox().enableAutoRange(x=False, y=False)
|
|
self.image = pg.ImageItem(axisOrder="row-major")
|
|
item.addItem(self.image)
|
|
self.history_by_combo: dict[tuple[int, int], deque[np.ndarray]] = {}
|
|
self.depth_axis_by_combo: dict[tuple[int, int], np.ndarray] = {}
|
|
self.signature: tuple[object, ...] | None = None
|
|
|
|
def sync(self, results: list[ResultCollection], window: int) -> bool:
|
|
# Completion time distinguishes recomputed payloads for the same acquisition.
|
|
signature = (window, tuple((collection_key(c), c.monotonic_ns, c.settings_revision) for c in results[-window:]))
|
|
if signature == self.signature:
|
|
return False
|
|
self.history_by_combo, self.depth_axis_by_combo = rebuild_bscan_history_from_results(results, window)
|
|
self.signature = signature
|
|
return True
|
|
|
|
def draw(self, key: tuple[int, int] | None, settings: GuiBscanStateModel) -> bool:
|
|
history = self.history_by_combo.get(key) if key is not None else None
|
|
depth = self.depth_axis_by_combo.get(key) if key is not None else None
|
|
if not history or depth is None:
|
|
self.image.clear()
|
|
return False
|
|
sweeps = apply_mean_ascan_subtraction(history, enabled=settings.subtract_mean_ascan)
|
|
depth_min, depth_max = float(np.min(depth)), float(np.max(depth))
|
|
count = sweeps.shape[0]
|
|
self.image.setImage(sweeps.T, autoLevels=False)
|
|
self.image.setRect(QRectF(.5, depth_min, float(count), max(depth_max - depth_min, 1e-6)))
|
|
self.image.setLookupTable(bscan_lookup_table(settings.axis))
|
|
self.image.setLevels(bscan_levels(sweeps, settings.axis))
|
|
item = self.widget.getPlotItem()
|
|
assert item is not None
|
|
item.getViewBox().setRange(xRange=(.5, .5 + count), yRange=(depth_min, depth_max), padding=.02)
|
|
assert key is not None
|
|
self.widget.setTitle(f"B-scan S21 in{key[0]}/out{key[1]} | sweeps={count}")
|
|
return True
|
|
|
|
def clear(self) -> None:
|
|
self.history_by_combo.clear()
|
|
self.depth_axis_by_combo.clear()
|
|
self.signature = None
|
|
self.image.clear()
|