ykcheungab
401 lines
13 KiB
Python
401 lines
13 KiB
Python
"""Qt-facing display backend for sensor viewer snapshots.
|
|
|
|
Separates acquisition (`gui.sensors.FrameSnapshot`) from Qt display state:
|
|
normalized images, force history / filtering, and aspect-ratio metadata.
|
|
Does not import camera, ROI, or `orisys.Sensor`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Literal, Optional, Sequence, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import orisys
|
|
from scipy.signal import butter, filtfilt
|
|
|
|
FORCE_SAMPLE_RATE_HZ = 30.0
|
|
PLOT_BUFFER_POINTS = 120
|
|
PLOT_WINDOW_SEC = PLOT_BUFFER_POINTS / FORCE_SAMPLE_RATE_HZ
|
|
Y_AXIS_MIN_SPAN = 20000.0
|
|
|
|
ArrowBackground = Literal["pure_color", "captured"]
|
|
# #0B1016 as BGR (matches demo_wuxi / promote_video theme)
|
|
PURE_COLOR_BG_BGR = (22, 16, 11)
|
|
|
|
_ARROW_DRAW_KW = dict(
|
|
threshold=2,
|
|
grid_spacing=20,
|
|
arrow_scale=0.5,
|
|
below_threshold_color=(160, 160, 160),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ViewerFrame:
|
|
"""Normalized display payload for one UI tick."""
|
|
|
|
display_image: np.ndarray
|
|
image_height: int
|
|
image_width: int
|
|
aspect_ratio: float # width / height
|
|
depth: Optional[np.ndarray]
|
|
normal_force: float
|
|
shear_x: float
|
|
shear_y: float
|
|
normal_force_calibrated: float
|
|
fps: float
|
|
timestamp: float
|
|
record_frame: Optional[np.ndarray]
|
|
|
|
|
|
@dataclass
|
|
class ForcePlotState:
|
|
"""Filtered force series ready for pyqtgraph curves."""
|
|
|
|
times: np.ndarray # shape (n,)
|
|
values: np.ndarray # shape (n, 3) FN, FX, FY
|
|
sample_count: int
|
|
x_range: Tuple[float, float]
|
|
y_ranges: Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]]
|
|
|
|
|
|
def lowpass_filter(
|
|
data,
|
|
cutoff_hz: float,
|
|
sample_rate_hz: float = FORCE_SAMPLE_RATE_HZ,
|
|
order: int = 2,
|
|
) -> np.ndarray:
|
|
"""Zero-phase Butterworth low-pass for uniformly sampled 1-D data."""
|
|
y = np.asarray(data, dtype=np.float64)
|
|
if cutoff_hz == 0:
|
|
return y.copy()
|
|
if y.size == 0:
|
|
return y
|
|
|
|
nyquist = 0.5 * sample_rate_hz
|
|
if cutoff_hz < 0 or cutoff_hz >= nyquist:
|
|
raise ValueError(
|
|
f"cutoff_hz must be in (0, {nyquist}) for sample_rate_hz={sample_rate_hz}"
|
|
)
|
|
|
|
b, a = butter(order, cutoff_hz / nyquist, btype="low")
|
|
padlen = 3 * (max(len(b), len(a)) - 1)
|
|
if y.size <= padlen:
|
|
return y.copy()
|
|
return filtfilt(b, a, y)
|
|
|
|
|
|
def ensure_bgr(img: Optional[np.ndarray]) -> Optional[np.ndarray]:
|
|
if img is None:
|
|
return None
|
|
if img.ndim == 2:
|
|
return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
|
if img.ndim == 3 and img.shape[2] == 1:
|
|
return cv2.cvtColor(img[:, :, 0], cv2.COLOR_GRAY2BGR)
|
|
return img
|
|
|
|
|
|
def image_shape_meta(img: np.ndarray) -> Tuple[int, int, float]:
|
|
"""Return (height, width, width/height) from a display image."""
|
|
h, w = int(img.shape[0]), int(img.shape[1])
|
|
aspect = float(w) / float(h) if h > 0 else 1.0
|
|
return h, w, aspect
|
|
|
|
|
|
def _base_bgr(snapshot) -> Optional[np.ndarray]:
|
|
for attr in ("image", "processed_image", "raw_image"):
|
|
out = ensure_bgr(getattr(snapshot, attr, None))
|
|
if out is not None:
|
|
return out
|
|
return None
|
|
|
|
|
|
def select_display_image(snapshot) -> np.ndarray:
|
|
"""Pick flow overlay / processed / raw image for the left panel."""
|
|
arrows = getattr(snapshot, "flow_overlay", None)
|
|
if arrows is not None:
|
|
out = ensure_bgr(arrows)
|
|
if out is not None:
|
|
return out
|
|
|
|
for attr in ("processed_image", "image", "raw_image"):
|
|
base = getattr(snapshot, attr, None)
|
|
out = ensure_bgr(base)
|
|
if out is not None:
|
|
return out
|
|
|
|
return np.zeros((400, 400, 3), dtype=np.uint8)
|
|
|
|
|
|
def build_display_image(
|
|
snapshot,
|
|
background: ArrowBackground = "pure_color",
|
|
) -> np.ndarray:
|
|
"""Build left-panel image: arrows on pure color or on the captured frame."""
|
|
flow = getattr(snapshot, "flow", None)
|
|
base = _base_bgr(snapshot)
|
|
|
|
if flow is not None and base is not None:
|
|
if background == "pure_color":
|
|
h, w = base.shape[:2]
|
|
canvas = np.empty((h, w, 3), dtype=np.uint8)
|
|
canvas[:, :] = PURE_COLOR_BG_BGR
|
|
return orisys.util.draw_arrows(canvas, flow, **_ARROW_DRAW_KW)
|
|
|
|
overlay = ensure_bgr(getattr(snapshot, "flow_overlay", None))
|
|
if overlay is not None:
|
|
return overlay
|
|
return orisys.util.draw_arrows(base, flow, **_ARROW_DRAW_KW)
|
|
|
|
return select_display_image(snapshot)
|
|
|
|
|
|
def bgr_to_rgb_bytes(img_bgr: np.ndarray) -> Tuple[np.ndarray, int, int]:
|
|
"""Convert BGR uint8 image to contiguous RGB for QImage."""
|
|
rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
if not rgb.flags["C_CONTIGUOUS"]:
|
|
rgb = np.ascontiguousarray(rgb)
|
|
h, w = rgb.shape[:2]
|
|
return rgb, h, w
|
|
|
|
|
|
def scale_bgr_to_fit(
|
|
img_bgr: np.ndarray,
|
|
max_w: int,
|
|
max_h: int,
|
|
) -> np.ndarray:
|
|
"""Resize preserving aspect ratio to fit inside (max_w, max_h)."""
|
|
if max_w <= 0 or max_h <= 0:
|
|
return img_bgr
|
|
h, w = img_bgr.shape[:2]
|
|
if h <= 0 or w <= 0:
|
|
return img_bgr
|
|
scale = min(max_w / float(w), max_h / float(h))
|
|
if abs(scale - 1.0) < 1e-6:
|
|
return img_bgr
|
|
nw = max(1, int(round(w * scale)))
|
|
nh = max(1, int(round(h * scale)))
|
|
return cv2.resize(img_bgr, (nw, nh), interpolation=cv2.INTER_AREA)
|
|
|
|
|
|
def panel_width_for_aspect(
|
|
available_height: int,
|
|
aspect_ratio: float,
|
|
*,
|
|
min_size: int = 160,
|
|
) -> int:
|
|
"""Preferred left-panel width given available height and image aspect."""
|
|
h = max(int(available_height), min_size)
|
|
w = int(round(h * float(aspect_ratio)))
|
|
return max(w, min_size)
|
|
|
|
|
|
class ForceHistory:
|
|
"""Ring-style force history buffer with filtering and y-range smoothing."""
|
|
|
|
def __init__(
|
|
self,
|
|
npoints: int = PLOT_BUFFER_POINTS,
|
|
*,
|
|
sample_rate_hz: float = FORCE_SAMPLE_RATE_HZ,
|
|
lowpass_cutoff_hz: float = 10.0,
|
|
y_axis_min_span: float = Y_AXIS_MIN_SPAN,
|
|
y_range_alpha: float = 0.5,
|
|
):
|
|
self.npoints = int(npoints)
|
|
self.sample_rate_hz = float(sample_rate_hz)
|
|
self.lowpass_cutoff_hz = float(lowpass_cutoff_hz)
|
|
self.y_axis_min_span = float(y_axis_min_span)
|
|
self.y_range_alpha = float(y_range_alpha)
|
|
self.window_sec = self.npoints / self.sample_rate_hz
|
|
self.reset()
|
|
|
|
def reset(self) -> None:
|
|
self.data = np.full((self.npoints, 2, 3), np.nan, dtype=np.float64)
|
|
self.sample_count = 0
|
|
self._y_range_smooth = [None, None, None]
|
|
|
|
def push(self, t: float, fn: float, fx: float, fy: float) -> None:
|
|
values = (float(fn), float(fx), float(fy))
|
|
if self.sample_count < self.npoints:
|
|
idx = self.sample_count
|
|
self.data[idx, 0, :] = t
|
|
self.data[idx, 1, :] = values
|
|
self.sample_count += 1
|
|
return
|
|
|
|
self.data[:-1, 0, :] = self.data[1:, 0, :]
|
|
self.data[:-1, 1, :] = self.data[1:, 1, :]
|
|
self.data[-1, 0, :] = t
|
|
self.data[-1, 1, :] = values
|
|
|
|
def _smooth_y_range(self, channel: int, ymin: float, ymax: float) -> Tuple[float, float]:
|
|
if self._y_range_smooth[channel] is None:
|
|
smoothed = (ymin, ymax)
|
|
else:
|
|
prev_min, prev_max = self._y_range_smooth[channel]
|
|
alpha = self.y_range_alpha
|
|
ymin = alpha * ymin + (1.0 - alpha) * prev_min
|
|
ymax = alpha * ymax + (1.0 - alpha) * prev_max
|
|
smoothed = (ymin, ymax)
|
|
|
|
span = smoothed[1] - smoothed[0]
|
|
if span < self.y_axis_min_span:
|
|
mid = 0.5 * (smoothed[0] + smoothed[1])
|
|
half = self.y_axis_min_span / 2.0
|
|
smoothed = (mid - half, mid + half)
|
|
|
|
self._y_range_smooth[channel] = smoothed
|
|
pad = max((smoothed[1] - smoothed[0]) * 0.12, 0.01)
|
|
return smoothed[0] - pad, smoothed[1] + pad
|
|
|
|
def plot_state(self) -> Optional[ForcePlotState]:
|
|
n = self.sample_count
|
|
if n == 0:
|
|
return None
|
|
|
|
data_plot = self.data[:n].copy()
|
|
for i in range(3):
|
|
data_plot[:, 1, i] = lowpass_filter(
|
|
data_plot[:, 1, i],
|
|
self.lowpass_cutoff_hz,
|
|
sample_rate_hz=self.sample_rate_hz,
|
|
)
|
|
|
|
times = data_plot[:, 0, 0].copy()
|
|
values = data_plot[:, 1, :].copy()
|
|
y_ranges = tuple(
|
|
self._smooth_y_range(i, float(values[:, i].min()), float(values[:, i].max()))
|
|
for i in range(3)
|
|
)
|
|
|
|
if n < self.npoints:
|
|
x_range = (0.0, self.window_sec)
|
|
else:
|
|
x0 = float(times[0])
|
|
x1 = float(times[-1])
|
|
if x1 <= x0:
|
|
x1 = x0 + self.window_sec
|
|
x_range = (x0, x1)
|
|
|
|
return ForcePlotState(
|
|
times=times,
|
|
values=values,
|
|
sample_count=n,
|
|
x_range=x_range,
|
|
y_ranges=y_ranges, # type: ignore[arg-type]
|
|
)
|
|
|
|
|
|
class QtViewerBackend:
|
|
"""Normalize FrameSnapshot into display images + force history."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
npoints: int = PLOT_BUFFER_POINTS,
|
|
lowpass_cutoff_hz: float = 10.0,
|
|
y_axis_min_span: float = Y_AXIS_MIN_SPAN,
|
|
arrow_background: ArrowBackground = "pure_color",
|
|
):
|
|
self.force_history = ForceHistory(
|
|
npoints=npoints,
|
|
lowpass_cutoff_hz=lowpass_cutoff_hz,
|
|
y_axis_min_span=y_axis_min_span,
|
|
)
|
|
self._start_time: Optional[float] = None
|
|
self.last_viewer_frame: Optional[ViewerFrame] = None
|
|
self.arrow_background: ArrowBackground = arrow_background
|
|
|
|
def set_arrow_background(self, background: ArrowBackground) -> None:
|
|
self.arrow_background = background
|
|
|
|
def rebuild_display(self, snapshot) -> ViewerFrame:
|
|
"""Re-render display image without advancing force history."""
|
|
display = build_display_image(snapshot, self.arrow_background)
|
|
h, w, aspect = image_shape_meta(display)
|
|
prev = self.last_viewer_frame
|
|
if prev is not None:
|
|
frame = ViewerFrame(
|
|
display_image=display,
|
|
image_height=h,
|
|
image_width=w,
|
|
aspect_ratio=aspect,
|
|
depth=prev.depth,
|
|
normal_force=prev.normal_force,
|
|
shear_x=prev.shear_x,
|
|
shear_y=prev.shear_y,
|
|
normal_force_calibrated=prev.normal_force_calibrated,
|
|
fps=prev.fps,
|
|
timestamp=prev.timestamp,
|
|
record_frame=prev.record_frame,
|
|
)
|
|
else:
|
|
frame = ViewerFrame(
|
|
display_image=display,
|
|
image_height=h,
|
|
image_width=w,
|
|
aspect_ratio=aspect,
|
|
depth=getattr(snapshot, "depth", None),
|
|
normal_force=float(getattr(snapshot, "normal_force", 0.0)),
|
|
shear_x=float(getattr(snapshot, "shear_x", 0.0)),
|
|
shear_y=float(getattr(snapshot, "shear_y", 0.0)),
|
|
normal_force_calibrated=float(
|
|
getattr(snapshot, "normal_force_calibrated", 0.0)
|
|
),
|
|
fps=float(getattr(snapshot, "fps", 0.0)),
|
|
timestamp=float(getattr(snapshot, "timestamp", 0.0) or 0.0),
|
|
record_frame=ensure_bgr(getattr(snapshot, "record_frame", None)),
|
|
)
|
|
self.last_viewer_frame = frame
|
|
return frame
|
|
|
|
def reset(self) -> None:
|
|
self.force_history.reset()
|
|
self._start_time = None
|
|
self.last_viewer_frame = None
|
|
|
|
def ingest(self, snapshot, *, now: Optional[float] = None) -> ViewerFrame:
|
|
"""Convert a FrameSnapshot into ViewerFrame and update force history."""
|
|
display = build_display_image(snapshot, self.arrow_background)
|
|
h, w, aspect = image_shape_meta(display)
|
|
|
|
fn = float(getattr(snapshot, "normal_force", 0.0))
|
|
fx = float(getattr(snapshot, "shear_x", 0.0))
|
|
fy = float(getattr(snapshot, "shear_y", 0.0))
|
|
fn_cal = float(getattr(snapshot, "normal_force_calibrated", 0.0))
|
|
fps = float(getattr(snapshot, "fps", 0.0))
|
|
ts = float(getattr(snapshot, "timestamp", 0.0) or 0.0)
|
|
if now is None:
|
|
now = ts if ts > 0 else 0.0
|
|
|
|
if self._start_time is None:
|
|
self._start_time = now
|
|
t = now - self._start_time
|
|
self.force_history.push(t, fn, fx, fy)
|
|
|
|
record = ensure_bgr(getattr(snapshot, "record_frame", None))
|
|
depth = getattr(snapshot, "depth", None)
|
|
|
|
frame = ViewerFrame(
|
|
display_image=display,
|
|
image_height=h,
|
|
image_width=w,
|
|
aspect_ratio=aspect,
|
|
depth=depth,
|
|
normal_force=fn,
|
|
shear_x=fx,
|
|
shear_y=fy,
|
|
normal_force_calibrated=fn_cal,
|
|
fps=fps,
|
|
timestamp=ts,
|
|
record_frame=record,
|
|
)
|
|
self.last_viewer_frame = frame
|
|
return frame
|
|
|
|
def plot_state(self) -> Optional[ForcePlotState]:
|
|
return self.force_history.plot_state()
|