orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add config dictionary to get_depth.py

Browse Source
This commit is contained in:
unknown
2025-12-30 10:25:16 +08:00
parent b9573c5781
commit 93e2984e7f
6 changed files with 358 additions and 334 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/orisys/get_depth.py
+337 -316
View File
@@ -12,325 +12,346 @@ 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
class Contact3D:
def __init__(self, config_contact=None, verbose=False):
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)。
Args:
flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
downsample_factor: 下采样因子，用于平滑计算（默认10）
return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
kernel_size: 形态学操作的核大小
denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
Returns:
dict: 包含扩张焦点信息的字典
"""
H, W = flow_field.shape[:2]
# 优化：使用视图而不是复制，减少内存分配
u = flow_field[:, :, 0]
v = flow_field[:, :, 1]
# 下采样并计算散度，然后插值回原尺寸
if downsample_factor > 1:
# 使用步进切片，避免复制
u_ds = u[::downsample_factor, ::downsample_factor]
v_ds = v[::downsample_factor, ::downsample_factor]
# 计算幅度（优化：使用square减少中间数组）
magnitude_ds = np.sqrt(np.square(u_ds) + np.square(v_ds))
if denoise_sigma > 0:
magnitude_ds = cv2.GaussianBlur(magnitude_ds, (0, 0), denoise_sigma)
# 优化：合并梯度计算，减少中间数组
# 计算散度：du/dx + dv/dy
du_dx = np.gradient(u_ds, axis=1)
dv_dy = np.gradient(v_ds, axis=0)
divergence_ds = du_dx + dv_dy # 直接相加，避免中间数组
if denoise_sigma > 0:
divergence_ds = cv2.GaussianBlur(divergence_ds, (0, 0), denoise_sigma)
# 插值回原尺寸
magnitude = cv2.resize(magnitude_ds, (W, H), interpolation=cv2.INTER_CUBIC)
divergence = cv2.resize(divergence_ds, (W, H), interpolation=cv2.INTER_CUBIC)
else:
# 计算幅度（优化：使用square减少中间数组）
magnitude = np.sqrt(np.square(u) + np.square(v))
if denoise_sigma > 0:
magnitude = cv2.GaussianBlur(magnitude, (0, 0), denoise_sigma)
# 优化：合并梯度计算
du_dx = np.gradient(u, axis=1)
dv_dy = np.gradient(v, axis=0)
divergence = du_dx + dv_dy
# 降噪：在计算散度之后对散度图进行平滑
if denoise_sigma > 0:
divergence = cv2.GaussianBlur(divergence, (0, 0), denoise_sigma)
# 优化：使用原地操作
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 = _create_foe_result(expansion_mask, contraction_mask, divergence)
# 寻找扩张区域轮廓
# 优化：缓存形态学操作的核
kernel_key = kernel_size
if kernel_key not in _morphology_kernels:
_morphology_kernels[kernel_key] = np.ones(kernel_size, np.uint8)
expansion_binary = expansion_mask.astype(np.uint8) * 255
if np.any(expansion_binary):
expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, _morphology_kernels[kernel_key])
contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
result['centroid_found'] = True
if return_all_contours:
result['_all_contours'] = contours
largest_contour = max(contours, key=cv2.contourArea)
result['contour'] = largest_contour
M = cv2.moments(largest_contour)
if M["m00"] > 0:
result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
return result
# Load Config
self.config_contact = config_contact
self.flow_threshold = config_contact["flow_threshold"]
self.expansion_threshold = config_contact["expansion_threshold"]
self.downsample_factor = config_contact["downsample_factor"]
self.kernel_size = config_contact["kernel_size"]
self.denoise_sigma = config_contact["denoise_sigma"]
self.min_contour_area = config_contact["min_contour_area"]
# Divergence threshold: only keep divergence values above this threshold, set others to zero
self.divergence_threshold = config_contact["div_threshold"]
self.verbose = verbose
if self.verbose:
print(f"Contact3D initialized with config: {self.config_contact}")
def compute_foe_gpu(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample_factor=1,
return_all_contours=False, kernel_size=(5,5), denoise_sigma=10):
"""
使用CUDA加速计算光流场的扩张焦点(Focus of Expansion)。
Args:
flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
downsample_factor: 下采样因子，用于平滑计算（默认10）
return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
kernel_size: 形态学操作的核大小
denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
Returns:
dict: 包含扩张焦点信息的字典
Note:
如果CuPy不可用，会自动回退到CPU版本的compute_foe
"""
# 检查CuPy是否可用
if not CUPY_AVAILABLE:
return compute_foe(flow_field, min_magnitude, expansion_threshold, downsample_factor,
return_all_contours, kernel_size, denoise_sigma)
H, W = flow_field.shape[:2]
# 导入CuPy的高斯滤波函数（仅在需要时导入）
try:
from cupyx.scipy.ndimage import gaussian_filter as cp_gaussian_filter
except ImportError:
# 如果导入失败，回退到CPU版本
return compute_foe(flow_field, min_magnitude, expansion_threshold, downsample_factor,
return_all_contours, kernel_size, denoise_sigma)
# 优化：使用asarray而不是array，如果已经是CuPy数组则避免复制
if isinstance(flow_field, cp.ndarray):
flow_gpu = flow_field
else:
flow_gpu = cp.asarray(flow_field, dtype=cp.float32)
# 优化：使用视图而不是复制
u_gpu = flow_gpu[:, :, 0]
v_gpu = flow_gpu[:, :, 1]
# 下采样并计算散度，然后插值回原尺寸
if downsample_factor > 1:
# 下采样（使用步进切片，避免复制）
u_ds_gpu = u_gpu[::downsample_factor, ::downsample_factor]
v_ds_gpu = v_gpu[::downsample_factor, ::downsample_factor]
# 计算幅度（优化：使用square减少中间数组）
magnitude_ds_gpu = cp.sqrt(cp.square(u_ds_gpu) + cp.square(v_ds_gpu))
# 高斯滤波（使用CuPy的gaussian_filter）
if denoise_sigma > 0:
magnitude_ds_gpu = cp_gaussian_filter(magnitude_ds_gpu, sigma=denoise_sigma)
# 优化：合并梯度计算，减少中间数组
# 计算散度：du/dx + dv/dy
du_dx_ds = cp.gradient(u_ds_gpu, axis=1)
dv_dy_ds = cp.gradient(v_ds_gpu, axis=0)
divergence_ds_gpu = du_dx_ds + dv_dy_ds
# 对散度进行高斯滤波
if denoise_sigma > 0:
divergence_ds_gpu = cp_gaussian_filter(divergence_ds_gpu, sigma=denoise_sigma)
# 优化：在GPU上完成插值，减少数据传输
# 使用CuPy的zoom进行插值（更高效）
from cupyx.scipy.ndimage import zoom as cp_zoom
zoom_factor = (H / divergence_ds_gpu.shape[0], W / divergence_ds_gpu.shape[1])
magnitude_gpu = cp_zoom(magnitude_ds_gpu, zoom_factor, order=3) # 3 = cubic
divergence_gpu = cp_zoom(divergence_ds_gpu, zoom_factor, order=3)
else:
# 计算幅度（优化：使用square减少中间数组）
magnitude_gpu = cp.sqrt(cp.square(u_gpu) + cp.square(v_gpu))
# 高斯滤波
if denoise_sigma > 0:
magnitude_gpu = cp_gaussian_filter(magnitude_gpu, sigma=denoise_sigma)
# 优化：合并梯度计算
du_dx = cp.gradient(u_gpu, axis=1)
dv_dy = cp.gradient(v_gpu, axis=0)
divergence_gpu = du_dx + dv_dy
# 对散度进行高斯滤波
if denoise_sigma > 0:
divergence_gpu = cp_gaussian_filter(divergence_gpu, sigma=denoise_sigma)
# 优化：在GPU上完成掩码操作，减少数据传输
# CuPy的where不支持out参数，使用赋值代替
divergence_gpu = cp.where(divergence_gpu > min_magnitude, divergence_gpu, 0)
# 在GPU上创建掩码
valid_mask_gpu = magnitude_gpu > min_magnitude
expansion_mask_gpu = (valid_mask_gpu & (divergence_gpu > expansion_threshold))
contraction_mask_gpu = (valid_mask_gpu & (divergence_gpu < -expansion_threshold))
# 优化：一次性传输所有需要的数据回CPU
magnitude = cp.asnumpy(magnitude_gpu)
divergence = cp.asnumpy(divergence_gpu)
expansion_mask = cp.asnumpy(expansion_mask_gpu)
contraction_mask = cp.asnumpy(contraction_mask_gpu)
# 初始化结果字典
result = _create_foe_result(expansion_mask, contraction_mask, divergence)
# 寻找扩张区域轮廓（在CPU上进行，因为OpenCV的轮廓检测需要CPU）
# 优化：缓存形态学操作的核
kernel_key = kernel_size
if kernel_key not in _morphology_kernels:
_morphology_kernels[kernel_key] = np.ones(kernel_size, np.uint8)
expansion_binary = expansion_mask.astype(np.uint8) * 255
if np.any(expansion_binary):
expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, _morphology_kernels[kernel_key])
contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
result['centroid_found'] = True
if return_all_contours:
result['_all_contours'] = contours
largest_contour = max(contours, key=cv2.contourArea)
result['contour'] = largest_contour
M = cv2.moments(largest_contour)
if M["m00"] > 0:
result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
return result
# Data Variables
self.foe_results = None
# self.expansion_mask = None
# self.contraction_mask = None
# self.divergence_map = None
# self.centroid = None
# self.centroid_found = False
# self.contour = None
# self.all_contours = None
def compute_multiple_foe(flow_field, min_magnitude=1.0, expansion_threshold=0.001,
downsample_factor=1, min_contour_area=3000, denoise_sigma=15, kernel_size=(5,5)):
"""
计算光流场的多个扩张焦点(Foci of Expansion)。
Args:
flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
downsample_factor: 下采样因子，用于平滑计算（默认10）
min_contour_area: 最小轮廓面积阈值，小于此值的轮廓会被忽略
denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
kernel_size: 形态学操作的核大小
def _create_foe_result(self, expansion_mask, contraction_mask, divergence_map,
centroid=None, centroid_found=False, contour=None,
all_contours=None):
"""
创建标准化的FOE结果字典，统一所有函数的返回格式。
Returns:
dict: 包含多个扩张焦点信息的字典，包括：
- 'foe_points': 扩张焦点坐标列表 [(x1, y1), (x2, y2), ...]
- 'expansion_regions': 扩张区域轮廓列表
- 'expansion_mask': 扩张区域二值掩码
- 'contraction_mask': 收缩区域二值掩码
- 'divergence_map': 散度图
"""
# 调用 compute_foe 获取基础数据和所有轮廓
foe_result = compute_foe_gpu(
flow_field,
min_magnitude=min_magnitude,
expansion_threshold=expansion_threshold,
downsample_factor=downsample_factor,
return_all_contours=True,
kernel_size=kernel_size,
denoise_sigma=denoise_sigma
)
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
# 构建多焦点结果，复用基础数据
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:
for contour in foe_result['_all_contours']:
# 过滤掉太小的轮廓
if cv2.contourArea(contour) >= min_contour_area:
# 计算轮廓的质心作为焦点
M = cv2.moments(contour)
def compute_foe(self, flow_field, return_all_contours=False):
# def compute_foe(self, 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)。
Args:
flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
downsample_factor: 下采样因子，用于平滑计算（默认10）
return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
kernel_size: 形态学操作的核大小
denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
Returns:
dict: 包含扩张焦点信息的字典
"""
H, W = flow_field.shape[:2]
# 优化：使用视图而不是复制，减少内存分配
u = flow_field[:, :, 0]
v = flow_field[:, :, 1]
# 下采样并计算散度，然后插值回原尺寸
if self.downsample_factor > 1:
# 使用步进切片，避免复制
u_ds = u[::self.downsample_factor, ::self.downsample_factor]
v_ds = v[::self.downsample_factor, ::self.downsample_factor]
# 计算幅度（优化：使用square减少中间数组）
magnitude_ds = np.sqrt(np.square(u_ds) + np.square(v_ds))
if self.denoise_sigma > 0:
magnitude_ds = cv2.GaussianBlur(magnitude_ds, (0, 0), self.denoise_sigma)
# 优化：合并梯度计算，减少中间数组
# 计算散度：du/dx + dv/dy
du_dx = np.gradient(u_ds, axis=1)
dv_dy = np.gradient(v_ds, axis=0)
divergence_ds = du_dx + dv_dy # 直接相加，避免中间数组
if self.denoise_sigma > 0:
divergence_ds = cv2.GaussianBlur(divergence_ds, (0, 0), self.denoise_sigma)
# 插值回原尺寸
magnitude = cv2.resize(magnitude_ds, (W, H), interpolation=cv2.INTER_CUBIC)
divergence = cv2.resize(divergence_ds, (W, H), interpolation=cv2.INTER_CUBIC)
else:
# 计算幅度（优化：使用square减少中间数组）
magnitude = np.sqrt(np.square(u) + np.square(v))
if self.denoise_sigma > 0:
magnitude = cv2.GaussianBlur(magnitude, (0, 0), self.denoise_sigma)
# 优化：合并梯度计算
du_dx = np.gradient(u, axis=1)
dv_dy = np.gradient(v, axis=0)
divergence = du_dx + dv_dy
# 降噪：在计算散度之后对散度图进行平滑
if self.denoise_sigma > 0:
divergence = cv2.GaussianBlur(divergence, (0, 0), self.denoise_sigma)
# 优化：使用原地操作 - 只保留大于阈值的散度值，其他设为0
divergence = np.where(divergence > self.divergence_threshold, divergence, 0)
# 创建掩码（使用原地操作优化）
valid_mask = magnitude > self.flow_threshold
expansion_mask = (valid_mask & (divergence > self.expansion_threshold))
contraction_mask = (valid_mask & (divergence < -self.expansion_threshold))
# 初始化结果字典
result = self._create_foe_result(expansion_mask, contraction_mask, divergence)
# 寻找扩张区域轮廓
# 优化：缓存形态学操作的核
# Convert kernel_size to tuple for use as dictionary key (lists are not hashable)
kernel_key = tuple(self.kernel_size) if isinstance(self.kernel_size, (list, np.ndarray)) else self.kernel_size
if kernel_key not in _morphology_kernels:
_morphology_kernels[kernel_key] = np.ones(self.kernel_size, np.uint8)
expansion_binary = expansion_mask.astype(np.uint8) * 255
if np.any(expansion_binary):
expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, _morphology_kernels[kernel_key])
contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
result['centroid_found'] = True
if return_all_contours:
result['_all_contours'] = contours
largest_contour = max(contours, key=cv2.contourArea)
result['contour'] = largest_contour
M = cv2.moments(largest_contour)
if M["m00"] > 0:
cx = int(M["m10"] / M["m00"])
cy = int(M["m01"] / M["m00"])
result['centroid'].append((cx, cy))
if 'contour' not in result:
result['contour'] = []
result['contour'].append(contour)
return result
result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
self.foe_results = result
return self.foe_results
def compute_foe_gpu(self, flow_field, return_all_contours=False):
"""
使用CUDA加速计算光流场的扩张焦点(Focus of Expansion)。
Args:
flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
flow_threshold: 最小光流幅度阈值，低于此值的点不参与计算
expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
downsample_factor: 下采样因子，用于平滑计算（默认10）
return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
kernel_size: 形态学操作的核大小
denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
Returns:
dict: 包含扩张焦点信息的字典
Note:
如果CuPy不可用，会自动回退到CPU版本的compute_foe
"""
# 检查CuPy是否可用
if not CUPY_AVAILABLE:
return self.compute_foe(flow_field, return_all_contours)
H, W = flow_field.shape[:2]
# 导入CuPy的高斯滤波函数（仅在需要时导入）
try:
from cupyx.scipy.ndimage import gaussian_filter as cp_gaussian_filter
except ImportError:
# 如果导入失败，回退到CPU版本
return self.compute_foe(flow_field, return_all_contours)
# 优化：使用asarray而不是array，如果已经是CuPy数组则避免复制
if isinstance(flow_field, cp.ndarray):
flow_gpu = flow_field
else:
flow_gpu = cp.asarray(flow_field, dtype=cp.float32)
# 优化：使用视图而不是复制
u_gpu = flow_gpu[:, :, 0]
v_gpu = flow_gpu[:, :, 1]
# 下采样并计算散度，然后插值回原尺寸
if self.downsample_factor > 1:
# 下采样（使用步进切片，避免复制）
u_ds_gpu = u_gpu[::self.downsample_factor, ::self.downsample_factor]
v_ds_gpu = v_gpu[::self.downsample_factor, ::self.downsample_factor]
# 计算幅度（优化：使用square减少中间数组）
magnitude_ds_gpu = cp.sqrt(cp.square(u_ds_gpu) + cp.square(v_ds_gpu))
# 高斯滤波（使用CuPy的gaussian_filter）
if self.denoise_sigma > 0:
magnitude_ds_gpu = cp_gaussian_filter(magnitude_ds_gpu, sigma=self.denoise_sigma)
# 优化：合并梯度计算，减少中间数组
# 计算散度：du/dx + dv/dy
du_dx_ds = cp.gradient(u_ds_gpu, axis=1)
dv_dy_ds = cp.gradient(v_ds_gpu, axis=0)
divergence_ds_gpu = du_dx_ds + dv_dy_ds
# 对散度进行高斯滤波
if self.denoise_sigma > 0:
divergence_ds_gpu = cp_gaussian_filter(divergence_ds_gpu, sigma=self.denoise_sigma)
# 优化：在GPU上完成插值，减少数据传输
# 使用CuPy的zoom进行插值（更高效）
from cupyx.scipy.ndimage import zoom as cp_zoom
zoom_factor = (H / divergence_ds_gpu.shape[0], W / divergence_ds_gpu.shape[1])
magnitude_gpu = cp_zoom(magnitude_ds_gpu, zoom_factor, order=3) # 3 = cubic
divergence_gpu = cp_zoom(divergence_ds_gpu, zoom_factor, order=3)
else:
# 计算幅度（优化：使用square减少中间数组）
magnitude_gpu = cp.sqrt(cp.square(u_gpu) + cp.square(v_gpu))
# 高斯滤波
if self.denoise_sigma > 0:
magnitude_gpu = cp_gaussian_filter(magnitude_gpu, sigma=self.denoise_sigma)
# 优化：合并梯度计算
du_dx = cp.gradient(u_gpu, axis=1)
dv_dy = cp.gradient(v_gpu, axis=0)
divergence_gpu = du_dx + dv_dy
# 对散度进行高斯滤波
if self.denoise_sigma > 0:
divergence_gpu = cp_gaussian_filter(divergence_gpu, sigma=self.denoise_sigma)
# 优化：在GPU上完成掩码操作，减少数据传输
# 只保留大于阈值的散度值，其他设为0
divergence_gpu = cp.where(divergence_gpu > self.divergence_threshold, divergence_gpu, 0)
# 在GPU上创建掩码
valid_mask_gpu = magnitude_gpu > self.flow_threshold
expansion_mask_gpu = (valid_mask_gpu & (divergence_gpu > self.expansion_threshold))
contraction_mask_gpu = (valid_mask_gpu & (divergence_gpu < -self.expansion_threshold))
# 优化：一次性传输所有需要的数据回CPU
magnitude = cp.asnumpy(magnitude_gpu)
divergence = cp.asnumpy(divergence_gpu)
expansion_mask = cp.asnumpy(expansion_mask_gpu)
contraction_mask = cp.asnumpy(contraction_mask_gpu)
# 初始化结果字典
result = self._create_foe_result(expansion_mask, contraction_mask, divergence)
# 寻找扩张区域轮廓（在CPU上进行，因为OpenCV的轮廓检测需要CPU）
# 优化：缓存形态学操作的核
# Convert kernel_size to tuple for use as dictionary key (lists are not hashable)
kernel_key = tuple(self.kernel_size) if isinstance(self.kernel_size, (list, np.ndarray)) else self.kernel_size
if kernel_key not in _morphology_kernels:
_morphology_kernels[kernel_key] = np.ones(self.kernel_size, np.uint8)
expansion_binary = expansion_mask.astype(np.uint8) * 255
if np.any(expansion_binary):
expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, _morphology_kernels[kernel_key])
contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
result['centroid_found'] = True
if return_all_contours:
result['_all_contours'] = contours
largest_contour = max(contours, key=cv2.contourArea)
result['contour'] = largest_contour
M = cv2.moments(largest_contour)
if M["m00"] > 0:
result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
self.foe_results = result
return self.foe_results
def compute_multiple_foe(self,flow_field):
"""
计算光流场的多个扩张焦点(Foci of Expansion)。
Args:
flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
flow_threshold: 最小光流幅度阈值，低于此值的点不参与计算
expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
downsample_factor: 下采样因子，用于平滑计算（默认10）
min_contour_area: 最小轮廓面积阈值，小于此值的轮廓会被忽略
denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
kernel_size: 形态学操作的核大小
Returns:
dict: 包含多个扩张焦点信息的字典，包括：
- 'foe_points': 扩张焦点坐标列表 [(x1, y1), (x2, y2), ...]
- 'expansion_regions': 扩张区域轮廓列表
- 'expansion_mask': 扩张区域二值掩码
- 'contraction_mask': 收缩区域二值掩码
- 'divergence_map': 散度图
"""
# 调用 compute_foe_gpu 获取基础数据和所有轮廓
foe_result = self.compute_foe_gpu(flow_field, return_all_contours=True)
# 构建多焦点结果，复用基础数据
result = self._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:
for contour in foe_result['_all_contours']:
# 过滤掉太小的轮廓
if cv2.contourArea(contour) >= self.min_contour_area:
# 计算轮廓的质心作为焦点
M = cv2.moments(contour)
if M["m00"] > 0:
cx = int(M["m10"] / M["m00"])
cy = int(M["m01"] / M["m00"])
result['centroid'].append((cx, cy))
if 'contour' not in result:
result['contour'] = []
result['contour'].append(contour)
self.foe_results = result
return self.foe_results