add focus of expansion (foe) detection

2025-12-28 12:20:44 +08:00
parent 52a8bf58dc
commit 017561578d
3 changed files with 263 additions and 13 deletions
@@ -0,0 +1,154 @@
+import numpy as np
+import cv2
+
+def compute_foe(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample_factor=10, return_all_contours=False,kernel_size=(5,5)):
+    """
+    计算光流场的扩张焦点(Focus of Expansion)。
+    
+    Args:
+        flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
+        min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
+        expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
+        downsample_factor: 下采样因子，用于平滑计算（默认10）
+        return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
+        
+    Returns:
+        dict: 包含扩张焦点信息的字典，包括：
+            - 'foe_point': 扩张焦点坐标 (x, y)，如果找到单个焦点
+            - 'expansion_region': 扩张区域轮廓（最大轮廓）
+            - 'expansion_mask': 扩张区域二值掩码
+            - 'contraction_mask': 收缩区域二值掩码
+            - 'divergence_map': 散度图
+            - '_all_contours': 所有轮廓列表（仅当return_all_contours=True时）
+    """
+    H, W = flow_field.shape[:2]
+    
+    # 平滑处理：下采样光流场，然后计算并插值回原尺寸
+    if downsample_factor > 1:
+        # 下采样光流场
+        flow_field_downsampled = flow_field[::downsample_factor, ::downsample_factor, :]
+        h_ds, w_ds = flow_field_downsampled.shape[:2]
+        
+        # 在下采样的光流场上计算幅度和散度
+        u_ds = flow_field_downsampled[:, :, 0]
+        v_ds = flow_field_downsampled[:, :, 1]
+        
+        # 计算下采样后的光流幅度
+        magnitude_ds = np.sqrt(u_ds**2 + v_ds**2)
+        
+        # 计算下采样后的散度
+        du_dx_ds = np.gradient(u_ds, axis=1)
+        dv_dy_ds = np.gradient(v_ds, axis=0)
+        divergence_ds = du_dx_ds + dv_dy_ds
+        
+        # 插值回原尺寸
+        magnitude = cv2.resize(magnitude_ds, (W, H), interpolation=cv2.INTER_CUBIC)
+        divergence = cv2.resize(divergence_ds, (W, H), interpolation=cv2.INTER_CUBIC)
+    else:
+        # 不进行下采样，直接计算
+        # 计算光流幅度
+        magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2)
+        
+        # 计算散度 (divergence)
+        u = flow_field[:, :, 0]
+        v = flow_field[:, :, 1]
+        
+        # 使用中心差分计算散度
+        du_dx = np.gradient(u, axis=1)
+        dv_dy = np.gradient(v, axis=0)
+        divergence = du_dx + dv_dy
+    
+    # 创建有效点掩码（幅度足够大的点）
+    valid_mask = magnitude > min_magnitude
+    
+    # 创建扩张和收缩区域掩码
+    expansion_mask = np.zeros_like(divergence, dtype=bool)
+    contraction_mask = np.zeros_like(divergence, dtype=bool)
+    
+    expansion_mask[valid_mask] = divergence[valid_mask] > expansion_threshold
+    contraction_mask[valid_mask] = divergence[valid_mask] < -expansion_threshold
+    
+    result = {
+        'expansion_mask': expansion_mask,
+        'contraction_mask': contraction_mask,
+        'divergence_map': divergence
+    }
+    
+    # 寻找扩张区域轮廓
+    expansion_binary = expansion_mask.astype(np.uint8) * 255
+    if np.any(expansion_binary):
+        # 应用形态学闭运算来填充小孔和连接邻近区域
+        kernel = np.ones(kernel_size, np.uint8)
+        expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, kernel)
+        
+        contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if contours:
+            # 如果请求返回所有轮廓，保存它们
+            if return_all_contours:
+                result['_all_contours'] = contours
+            
+            # 找到最大的扩张区域轮廓
+            largest_contour = max(contours, key=cv2.contourArea)
+            result['expansion_region'] = largest_contour
+            
+            # 计算扩张区域的质心作为焦点
+            M = cv2.moments(largest_contour)
+            if M["m00"] > 0:
+                cx = int(M["m10"] / M["m00"])
+                cy = int(M["m01"] / M["m00"])
+                result['foe_point'] = (cx, cy)
+    
+    return result
+
+def compute_multiple_foe(flow_field, min_magnitude=1.0, expansion_threshold=0.001, 
+                         downsample_factor=10, min_contour_area=3000):
+    """
+    计算光流场的多个扩张焦点(Foci of Expansion)。
+    
+    Args:
+        flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
+        min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
+        expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
+        downsample_factor: 下采样因子，用于平滑计算（默认10）
+        min_contour_area: 最小轮廓面积阈值，小于此值的轮廓会被忽略
+        
+    Returns:
+        dict: 包含多个扩张焦点信息的字典，包括：
+            - 'foe_points': 扩张焦点坐标列表 [(x1, y1), (x2, y2), ...]
+            - 'expansion_regions': 扩张区域轮廓列表
+            - 'expansion_mask': 扩张区域二值掩码
+            - 'contraction_mask': 收缩区域二值掩码
+            - 'divergence_map': 散度图
+    """
+    # 调用 compute_foe 获取基础数据和所有轮廓
+    foe_result = compute_foe(
+        flow_field, 
+        min_magnitude=min_magnitude,
+        expansion_threshold=expansion_threshold,
+        downsample_factor=downsample_factor,
+        return_all_contours=True
+    )
+    
+    # 构建多焦点结果，复用基础数据
+    result = {
+        'expansion_mask': foe_result['expansion_mask'],
+        'contraction_mask': foe_result['contraction_mask'],
+        'divergence_map': foe_result['divergence_map'],
+        'foe_points': [],
+        'expansion_regions': []
+    }
+    
+    # 处理所有轮廓，过滤并计算焦点
+    if '_all_contours' in foe_result:
+        for contour in foe_result['_all_contours']:
+            # 过滤掉太小的轮廓
+            if cv2.contourArea(contour) >= min_contour_area:
+                # 计算轮廓的质心作为焦点
+                M = cv2.moments(contour)
+                if M["m00"] > 0:
+                    cx = int(M["m10"] / M["m00"])
+                    cy = int(M["m01"] / M["m00"])
+                    result['foe_points'].append((cx, cy))
+                    result['expansion_regions'].append(contour)
+    
+    return result