fixed new sample cannot detect boxes due to filter

2025-12-29 15:06:06 +08:00
parent 4f8165b04b
commit 55228413d4
10 changed files with 180 additions and 92 deletions
@@ -12,6 +12,41 @@ except ImportError:
 # 缓存形态学操作的核，避免重复创建
 _morphology_kernels = {}

+def _create_foe_result(expansion_mask, contraction_mask, divergence_map, 
+                       centroid=None, centroid_found=False, contour=None, 
+                       all_contours=None):
+    """
+    创建标准化的FOE结果字典，统一所有函数的返回格式。
+    
+    Args:
+        expansion_mask: 扩张区域掩码
+        contraction_mask: 收缩区域掩码
+        divergence_map: 散度图
+        centroid: 质心坐标，可以是单个元组 (x, y) 或列表 [(x1, y1), (x2, y2), ...]
+        centroid_found: 是否找到质心
+        contour: 轮廓，可以是单个轮廓或轮廓列表
+        all_contours: 所有轮廓（仅在return_all_contours=True时使用）
+    
+    Returns:
+        dict: 标准化的FOE结果字典
+    """
+    result = {
+        'expansion_mask': expansion_mask,
+        'contraction_mask': contraction_mask,
+        'divergence_map': divergence_map,
+        'centroid': centroid if centroid is not None else [0,0],
+        'centroid_found': centroid_found,
+    }
+    
+    # 可选字段
+    if contour is not None:
+        result['contour'] = contour
+    
+    if all_contours is not None:
+        result['_all_contours'] = all_contours
+    
+    return result
+
 def compute_foe(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample_factor=1, return_all_contours=False, kernel_size=(5,5), denoise_sigma=10):
    """
    计算光流场的扩张焦点(Focus of Expansion)。
@@ -73,20 +108,15 @@ def compute_foe(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample
            divergence = cv2.GaussianBlur(divergence, (0, 0), denoise_sigma)
    
    # 优化：使用原地操作
-    np.where(divergence > 0.02, divergence, 0, out=divergence)
+    np.where(divergence > min_magnitude, divergence, 0, out=divergence)
    
    # 创建掩码（使用原地操作优化）
    valid_mask = magnitude > min_magnitude
    expansion_mask = (valid_mask & (divergence > expansion_threshold))
    contraction_mask = (valid_mask & (divergence < -expansion_threshold))
    
-    result = {
-        'expansion_mask': expansion_mask,
-        'contraction_mask': contraction_mask,
-        'divergence_map': divergence,
-        'foe_points': [0, 0],
-        'is_foe_found': False
-    }
+    # 初始化结果字典
+    result = _create_foe_result(expansion_mask, contraction_mask, divergence)
    
    # 寻找扩张区域轮廓
    # 优化：缓存形态学操作的核
@@ -100,16 +130,16 @@ def compute_foe(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample
        contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
-            result['is_foe_found'] = True
+            result['centroid_found'] = True
            if return_all_contours:
                result['_all_contours'] = contours
            
            largest_contour = max(contours, key=cv2.contourArea)
-            result['expansion_region'] = largest_contour
+            result['contour'] = largest_contour
            
            M = cv2.moments(largest_contour)
            if M["m00"] > 0:
-                result['foe_point'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
+                result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
    
    return result

@@ -206,7 +236,7 @@ def compute_foe_gpu(flow_field, min_magnitude=1.0, expansion_threshold=1, downsa
    
    # 优化：在GPU上完成掩码操作，减少数据传输
    # CuPy的where不支持out参数，使用赋值代替
-    divergence_gpu = cp.where(divergence_gpu > 0.02, divergence_gpu, 0)
+    divergence_gpu = cp.where(divergence_gpu > min_magnitude, divergence_gpu, 0)
    
    # 在GPU上创建掩码
    valid_mask_gpu = magnitude_gpu > min_magnitude
@@ -219,13 +249,8 @@ def compute_foe_gpu(flow_field, min_magnitude=1.0, expansion_threshold=1, downsa
    expansion_mask = cp.asnumpy(expansion_mask_gpu)
    contraction_mask = cp.asnumpy(contraction_mask_gpu)
    
-    result = {
-        'expansion_mask': expansion_mask,
-        'contraction_mask': contraction_mask,
-        'divergence_map': divergence,
-        'foe_points': [0, 0],
-        'is_foe_found': False
-    }
+    # 初始化结果字典
+    result = _create_foe_result(expansion_mask, contraction_mask, divergence)
    
    # 寻找扩张区域轮廓（在CPU上进行，因为OpenCV的轮廓检测需要CPU）
    # 优化：缓存形态学操作的核
@@ -239,16 +264,16 @@ def compute_foe_gpu(flow_field, min_magnitude=1.0, expansion_threshold=1, downsa
        contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
-            result['is_foe_found'] = True
+            result['centroid_found'] = True
            if return_all_contours:
                result['_all_contours'] = contours
            
            largest_contour = max(contours, key=cv2.contourArea)
-            result['expansion_region'] = largest_contour
+            result['contour'] = largest_contour
            
            M = cv2.moments(largest_contour)
            if M["m00"] > 0:
-                result['foe_point'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
+                result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
    
    return result

@@ -275,7 +300,7 @@ def compute_multiple_foe(flow_field, min_magnitude=1.0, expansion_threshold=0.00
            - 'divergence_map': 散度图
    """
    # 调用 compute_foe 获取基础数据和所有轮廓
-    foe_result = compute_foe(
+    foe_result = compute_foe_gpu(
        flow_field, 
        min_magnitude=min_magnitude,
        expansion_threshold=expansion_threshold,
@@ -286,13 +311,12 @@ def compute_multiple_foe(flow_field, min_magnitude=1.0, expansion_threshold=0.00
    )
    
    # 构建多焦点结果，复用基础数据
-    result = {
-        'expansion_mask': foe_result['expansion_mask'],
-        'contraction_mask': foe_result['contraction_mask'],
-        'divergence_map': foe_result['divergence_map'],
-        'foe_points': [],
-        'expansion_regions': []
-    }
+    result = _create_foe_result(
+        expansion_mask=foe_result['expansion_mask'],
+        contraction_mask=foe_result['contraction_mask'],
+        divergence_map=foe_result['divergence_map'],
+        centroid_found=foe_result['centroid_found']
+    )
    
    # 处理所有轮廓，过滤并计算焦点
    if '_all_contours' in foe_result:
@@ -304,7 +328,9 @@ def compute_multiple_foe(flow_field, min_magnitude=1.0, expansion_threshold=0.00
                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)
+                    result['centroid'].append((cx, cy))
+                    if 'contour' not in result:
+                        result['contour'] = []
+                    result['contour'].append(contour)
    
    return result