orisys-sdk/demo.py

import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import numpy as np
from orisys import Sensor

def draw_arrows_fade(img, flow, grid_spacing=10, arrow_scale=1.0, threshold=5.0):
    """
    Draw optical flow arrows on image with transparency based on magnitude.
    
    Args:
        img: Input image (BGR or grayscale)
        flow: Optical flow field with shape (flow_height, flow_width, 2)
        grid_spacing: Spacing between arrows in pixels (default: 20)
        arrow_scale: Scale factor for arrow length (default: 1.0)
        threshold: Minimum flow magnitude to display arrow (default: 0.0)
    
    Returns:
        Image with arrows drawn, blended with original image
    """
    # Make a copy of the image
    img_copy = img.copy()
    
    # Convert grayscale to BGR if needed
    if len(img_copy.shape) == 2 or img_copy.shape[2] == 1:
        img_copy = cv2.cvtColor(img_copy, cv2.COLOR_GRAY2BGR)
    
    # Get image and flow dimensions
    img_height, img_width = img_copy.shape[:2]
    flow_height, flow_width = flow.shape[:2]
    
    # Calculate scaling factors between image and flow dimensions
    scale_x = img_width / flow_width
    scale_y = img_height / flow_height
    
    # Create grid positions in flow coordinates
    grid_y, grid_x = np.mgrid[grid_spacing//2:flow_height:grid_spacing, 
                                grid_spacing//2:flow_width:grid_spacing]
    
    # Flatten grid positions
    grid_y_flat = grid_y.flatten()
    grid_x_flat = grid_x.flatten()
    
    # Get flow vectors at grid positions
    flow_vectors = flow[grid_y_flat, grid_x_flat, :]  # Shape: (n_points, 2)
    
    # Calculate flow magnitudes
    magnitudes = np.linalg.norm(flow_vectors, axis=1)
    
    # Filter arrows by threshold: only show arrows above threshold
    valid_mask = magnitudes >= threshold
    
    # Normalize magnitudes to 0-1 range for alpha blending
    # Zero magnitude = transparent, higher magnitude = more opaque white
    max_magnitude = magnitudes.max() if magnitudes.max() > 0 else 1.0
    alpha_values = magnitudes / max_magnitude
    
    # Scale grid positions to image coordinates
    img_grid_x = (grid_x_flat * scale_x).astype(np.int32)
    img_grid_y = (grid_y_flat * scale_y).astype(np.int32)
    
    # Scale flow vectors to image coordinates
    scaled_flow_x = flow_vectors[:, 0] * scale_x / arrow_scale
    scaled_flow_y = flow_vectors[:, 1] * scale_y / arrow_scale
    
    # Calculate arrow endpoints
    end_x = (img_grid_x + scaled_flow_x).astype(np.int32)
    end_y = (img_grid_y + scaled_flow_y).astype(np.int32)
    
    # Create overlay image with white background for arrows
    overlay = np.zeros_like(img_copy)
    
    # Create weight map for blending based on arrow presence and magnitude
    weight_map = np.zeros((img_height, img_width), dtype=np.float32)
    
    # Draw arrows on overlay with white color and build weight map
    for i in range(len(img_grid_x)):
        # Skip if magnitude is below threshold
        if not valid_mask[i]:
            continue
            
        # Check bounds
        if (0 <= img_grid_x[i] < img_width and 0 <= img_grid_y[i] < img_height and
            0 <= end_x[i] < img_width and 0 <= end_y[i] < img_height):
            
            # Get alpha value (0 = transparent, 1 = opaque white)
            alpha = alpha_values[i]
            
            # Skip if magnitude is zero (transparent)
            if alpha <= 0:
                continue
            
            # Calculate arrow thickness based on magnitude
            thickness = max(1, int(1 + 2 * alpha))
            
            # White color for arrows
            color = (255, 255, 255)
            
            # Draw arrow on overlay
            cv2.arrowedLine(overlay, 
                          (img_grid_x[i], img_grid_y[i]), 
                          (end_x[i], end_y[i]), 
                          color, 
                          thickness=thickness, 
                          tipLength=0.3)
            
            # Create a line mask for this arrow in the weight map
            line_mask = np.zeros((img_height, img_width), dtype=np.float32)
            cv2.line(line_mask,
                    (img_grid_x[i], img_grid_y[i]),
                    (end_x[i], end_y[i]),
                    1.0,
                    thickness=thickness)
            # Add arrow tip approximation
            cv2.circle(line_mask, (end_x[i], end_y[i]), thickness, 1.0, -1)
            
            # Weight by alpha (magnitude) - use maximum for overlapping arrows
            weight_map = np.maximum(weight_map, line_mask * alpha)
    
    # Expand weight map to 3 channels
    weight_map_3ch = np.stack([weight_map] * 3, axis=2)
    
    # Blend: result = original * (1 - weight) + overlay * weight
    result = (img_copy.astype(np.float32) * (1 - weight_map_3ch) + 
              overlay.astype(np.float32) * weight_map_3ch)
    
    return result.astype(np.uint8)

def main():
    global start_time
    
    # 打开视频
    sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True)
    # 打开摄像头
    # sensor1 = Sensor(0,isstitch=True,cuda = True,verbose = True)
    while True:
        sensor1.get_img() # 提取并拼接复眼图像
        sensor1.compute_deformation() # 计算形变和分解
        # 提取信息：
        # 向量场: VRAW, VNORMAL, VSHEAR: 原始，法向，切向 
        # 帧率，时间戳: FPS, TIMESTAMP
        # 变形量总和: FNORMAL , FSHEARX, FSHEARY : 法向，切向（X)，切向（Y)
        fps, fn, fx, fy, flow = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW)
        arrows = draw_arrows_fade(sensor1.img, flow)
        cv2.imshow("Deformation", arrows)

        key = cv2.waitKey(1)
        if key & 0xFF == ord("q"):
            print("\n 停止程序")
            break
    
    sensor1.disconnect() # 断开传感器连接



if __name__ == '__main__':
    main()