added mesh surface arrows

src/orisys/util.py

@@ -1,16 +1,50 @@
 import cv2
 import numpy as np
 
-def draw_arrows(frame, optical_flow, fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0, margin = 0.8, colormap=cv2.COLORMAP_JET, below_threshold_color=(0, 0, 255), edge_margin=15):
+
+def _prepare_draw_mask(valid_mask, shape):
+    """Resize a boolean mask to (height, width) for arrow culling."""
+    if valid_mask is None:
+        return None
+    h, w = shape[:2]
+    mask = np.asarray(valid_mask, dtype=bool)
+    if mask.shape != (h, w):
+        mask_u8 = mask.astype(np.uint8)
+        if mask_u8.ndim > 2:
+            mask_u8 = mask_u8[:, :, 0]
+        mask = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
+    return mask
+
+
+def draw_arrows(
+    frame,
+    optical_flow,
+    fps=True,
+    threshold=5,
+    grid_spacing=40,
+    arrow_scale=1.0,
+    margin=0.8,
+    colormap=cv2.COLORMAP_JET,
+    below_threshold_color=(0, 0, 255),
+    edge_margin=15,
+    valid_mask=None,
+):
     """
     Draw colored optical flow arrows on the image by magnitude (JET colormap).
     Frame and optical_flow must have the same height/width; arrows are drawn in pixel coordinates.
+
+    Parameters
+    ----------
+    valid_mask : array-like, optional
+        Boolean (H, W) mask. Grid points outside the mask are skipped (no arrows or dots).
     """
     image = frame.copy()
     if len(image.shape) == 2 or image.shape[2] == 1:
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
 
     h, w = image.shape[:2]
+    mask = _prepare_draw_mask(valid_mask, (h, w))
+
     # Arrow grid in pixel coordinates (assume frame and flow same size)
     rr = np.arange(grid_spacing // 2, h, grid_spacing)
     cc = np.arange(grid_spacing // 2, w, grid_spacing)
@@ -24,17 +58,29 @@ def draw_arrows(frame, optical_flow, fps=True, threshold=5, grid_spacing=40, arr
     dy = disp[:, 1] / arrow_scale
     length = np.sqrt(dx * dx + dy * dy)
     # Cap arrow length at grid_spacing
-    scale = np.where(length > grid_spacing*margin, grid_spacing / np.maximum(length, 1e-9), 1.0)*margin
+    scale = np.where(length > grid_spacing * margin, grid_spacing / np.maximum(length, 1e-9), 1.0) * margin
     dx, dy = dx * scale, dy * scale
     end_x = (cc + dx).astype(np.int32)
     end_y = (rr + dy).astype(np.int32)
 
     valid = (
-        (cc >= edge_margin) & (cc < w - edge_margin) & (rr >= edge_margin) & (rr < h - edge_margin)
-        & (end_x >= 0) & (end_x < w) & (end_y >= 0) & (end_y < h)
+        (cc >= edge_margin)
+        & (cc < w - edge_margin)
+        & (rr >= edge_margin)
+        & (rr < h - edge_margin)
+        & (end_x >= 0)
+        & (end_x < w)
+        & (end_y >= 0)
+        & (end_y < h)
     )
+    if mask is not None:
+        valid &= mask[rr, cc]
+        valid &= mask[end_y, end_x]
+
     mag_range = mag.max() - mag.min()
-    norm_mag = ((mag - mag.min()) / mag_range * 255).astype(np.uint8) if mag_range > 0 else np.zeros_like(mag, dtype=np.uint8)
+    norm_mag = (
+        ((mag - mag.min()) / mag_range * 255).astype(np.uint8) if mag_range > 0 else np.zeros_like(mag, dtype=np.uint8)
+    )
     norm_mag[mag < threshold] = 0
     colors = cv2.applyColorMap(norm_mag.reshape(-1, 1), colormap)
 
@@ -44,7 +90,6 @@ def draw_arrows(frame, optical_flow, fps=True, threshold=5, grid_spacing=40, arr
         if mag[i] < threshold:
             cv2.circle(image, pt1, radius=1, color=below_threshold_color, thickness=-1)
         else:
-            # thickness = max(1, int(2 * mag[i] / mag.max())) if mag.max() > 0 else 2
             thickness = 1
             cv2.arrowedLine(image, pt1, pt2, color, thickness=thickness, tipLength=0.2)
     return image