split example visualizaiton pipeline

examples/viz/opencv_maps.py

@@ -0,0 +1,57 @@
+"""OpenCV visualization helpers for example scripts."""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+
+def depth_colormap(depth_map, colormap=cv2.COLORMAP_JET) -> np.ndarray:
+    """Normalize a scalar map to uint8 and apply an OpenCV colormap."""
+    values = np.asarray(depth_map)
+    if values.size == 0:
+        return np.zeros((1, 1, 3), dtype=np.uint8)
+    if values.max() > values.min():
+        normalized = ((values - values.min()) / (values.max() - values.min()) * 255).astype(np.uint8)
+    else:
+        normalized = np.zeros_like(values, dtype=np.uint8)
+    return cv2.applyColorMap(normalized, colormap)
+
+
+def draw_magnitude_map(
+    vfield,
+    threshold=3.0,
+    colormap=cv2.COLORMAP_JET,
+    output_size=(400, 400),
+) -> np.ndarray:
+    """Render a smoothed optical-flow magnitude map as a BGR pseudo-color image."""
+    magnitude = np.sqrt(vfield[:, :, 0] ** 2 + vfield[:, :, 1] ** 2)
+
+    downsample_factor = 10
+    h, w = magnitude.shape
+    magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
+    magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC)
+
+    magnitude_thresholded = magnitude_smoothed.copy()
+    magnitude_thresholded[magnitude_smoothed < threshold] = 0
+
+    mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max()
+    if mag_max > mag_min:
+        magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8)
+    else:
+        magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)
+
+    magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
+    out_w, out_h = output_size
+    return cv2.resize(magnitude_colored, (out_w, out_h), interpolation=cv2.INTER_CUBIC)
+
+
+def frame_to_bgr(frame: np.ndarray) -> np.ndarray:
+    """Convert a grayscale or single-channel frame to BGR for display/recording."""
+    if frame is None:
+        raise ValueError("frame is None")
+    if len(frame.shape) == 2:
+        return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
+    if len(frame.shape) == 3 and frame.shape[2] == 1:
+        return cv2.cvtColor(frame.squeeze(2), cv2.COLOR_GRAY2BGR)
+    return frame