create util.py

src/orisys/util.py

@@ -0,0 +1,153 @@
+import cv2
+import numpy as np
+from scipy.interpolate import griddata
+
+def draw_arrows(image, optical_flow,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0):
+    """
+    Draw colored optical flow arrows on sensor image based on flow magnitude.
+    
+    Args:
+        threshold: Minimum flow magnitude to display arrow (default: 10.0)
+        grid_spacing: Spacing between arrows in pixels (default: 20)
+        arrow_scale: Scale factor for arrow length (default: 3.0)
+    
+    Returns:
+        np.ndarray: 400x400 BGR image with colored arrows (JET colormap)
+    
+    """
+    # Create working copies to avoid modifying originals
+    
+    # Convert grayscale to BGR to enable colored arrow drawing
+    if len(image.shape) == 2 or image.shape[2] == 1:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+    
+    # Calculate scaling factors between image and flow dimensions
+    img_height, img_width = image.shape[:2]
+    flow_height, flow_width = optical_flow.shape[:2]
+    scale_x, scale_y = img_width / flow_width, img_height / flow_height
+    
+    # Create sampling grid for flow field (used for interpolation source points)
+    sample_step = 19  # Spacing between sample points
+    flow_start = np.stack(np.meshgrid(range(0, flow_width, sample_step), range(0, flow_height, sample_step)), 2)
+    flow_end = (optical_flow[flow_start[:, :, 1], flow_start[:, :, 0], :] * 2.0 + flow_start).astype(np.int32)
+    
+    # Create regular grid for arrow placement in image coordinates
+    grid_x, grid_y = np.mgrid[grid_spacing//2:flow_height:grid_spacing, grid_spacing//2:flow_width:grid_spacing]
+    x_axis = (grid_y.reshape(-1) * scale_x).astype(np.int64)  # Scale to image width
+    y_axis = (grid_x.reshape(-1) * scale_y).astype(np.int64)  # Scale to image height
+    
+    # Interpolate flow vectors from sample grid to regular arrow grid using cubic interpolation
+    grid_vector = griddata(
+        flow_start.reshape(-1, 2), 
+        (flow_end - flow_start).reshape(-1, 2),  # Flow from start to end (original to new position)
+        (grid_y, grid_x), method='cubic', fill_value=0
+    )
+    grid_vector[np.isnan(grid_vector)] = 0  # Replace NaN with zero for stability
+    grid_flat_scaled = grid_vector.reshape(-1, 2) * np.array([scale_x, scale_y])
+    
+    # Compute arrow endpoints and ensure they stay within image bounds
+    end_x = (grid_flat_scaled[:, 0] / arrow_scale + x_axis).astype(np.int64)
+    end_y = (grid_flat_scaled[:, 1] / arrow_scale + y_axis).astype(np.int64)
+    valid_mask = (y_axis >= 0) & (y_axis < img_height) & (x_axis >= 0) & (x_axis < img_width)
+    valid_mask &= (end_x >= 0) & (end_x < img_width) & (end_y >= 0) & (end_y < img_height)
+    
+    # Filter arrows by magnitude threshold and compute colors using JET colormap
+    magnitudes = np.linalg.norm(grid_flat_scaled, axis=1)
+    valid_mask &= (magnitudes >= 0)  # Only show arrows above threshold
+    valid_indices = np.where(valid_mask)[0]
+    mag_range = magnitudes.max() - magnitudes.min()
+    # Normalize magnitudes to 0-255 range for colormap application, but set below-threshold values to 0 (blue)
+    norm_magnitudes = ((magnitudes - magnitudes.min()) / mag_range * 255.0).astype(np.uint8) if mag_range > 0 else np.zeros_like(magnitudes, dtype=np.uint8)
+    norm_magnitudes[magnitudes < threshold] = 0  # Keep lowest color (blue) for below-threshold magnitudes
+    colors_array = cv2.applyColorMap(norm_magnitudes.reshape(-1, 1), cv2.COLORMAP_JET)  # JET: blue (low) -> red (high)
+    
+    # Draw arrows at each valid grid point with magnitude-based coloring
+    for i in valid_indices:
+        color = tuple(map(int, colors_array[i, 0]))
+        # Scale arrow length and thickness based on threshold
+        if magnitudes[i] < threshold:
+            # Below threshold: draw as small dot (zero-length arrow with larger tip)
+            cv2.circle(image, (x_axis[i], y_axis[i]), radius=1, color=color, thickness=-1)
+        else:
+            # Above threshold: draw full arrow with thickness proportional to magnitude
+            # Arrow rooted at original position (x_axis, y_axis), head points to new position (end_x, end_y)
+            thickness = max(1, int(2 * magnitudes[i] / magnitudes.max())) if magnitudes.max() > 0 else 2
+            cv2.arrowedLine(image, (x_axis[i], y_axis[i]), (end_x[i], end_y[i]), color, thickness=thickness, tipLength=0.3)
+    return image
+
+def draw_magnitude_map(image, flow_field, threshold=3.0, colormap=cv2.COLORMAP_JET):
+    """
+    Draw color map of pointwise magnitude of velocity field.
+    
+    Args:
+        info_type: Type of flow field to visualize (VRAW, VNORMAL, VSHEAR, or None for total)
+        threshold: Minimum magnitude to display (below threshold shows as zero/blue)
+        colormap: OpenCV colormap to use (default: COLORMAP_JET)
+    
+    Returns:
+        np.ndarray: 400x400 BGR image with colored magnitude map
+    """
+    # Get the appropriate flow field
+    # if info_type is None or info_type == self.info.VRAW:
+    #     flow_field = self.flow.copy() if self.flow is not None else None
+    # elif info_type == self.info.VNORMAL:
+    #     flow_field = self.normal.copy() if self.normal is not None else None
+    # elif info_type == self.info.VSHEAR:
+    #     flow_field = self.shear.copy() if self.shear is not None else None
+    # else:
+    #     # For combined normal + shear magnitude
+    #     if self.normal is not None and self.shear is not None:
+    #         flow_field = self.normal + self.shear
+    #     else:
+    #         flow_field = None
+    
+    # if flow_field is None:
+    #     # Return blank image if no flow data
+    #     return np.zeros((self.config_resolution_stitch[0], self.config_resolution_stitch[1], 3), dtype=np.uint8)
+    
+    # Compute magnitude at each point
+    magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2)
+    
+    # Smooth by downsampling and interpolating back
+    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)
+    
+    # Apply threshold - set below-threshold values to zero
+    magnitude_thresholded = magnitude_smoothed.copy()
+    magnitude_thresholded[magnitude_smoothed < threshold] = 0
+    
+    # Normalize to 0-255 range
+    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)
+    
+    # Apply colormap
+    magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
+    
+    # Resize to output size
+    # out_img = cv2.resize(magnitude_colored, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC)
+    
+    return magnitude_colored
+
+def draw_fps(image, fps):
+    """
+    Draw FPS counter on the top left corner of the image.
+    """
+    font_scale=0.5
+    thickness=2 
+    color=(0, 255, 0)
+    image = image.copy()
+    
+    # Draw FPS text
+    fps_text = f"FPS: {fps:.1f}"
+    position = (10, 30)  # Top left corner with margin
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    
+    # Draw FPS text
+    cv2.putText(image, fps_text, position, font, font_scale, color, thickness, cv2.LINE_AA)
+    
+    return image