added contact gpu support
This commit is contained in:
bluepulsar
+191
-14
@@ -1,6 +1,17 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
# Try to import CuPy for GPU acceleration
|
||||
try:
|
||||
import cupy as cp
|
||||
CUPY_AVAILABLE = True
|
||||
except ImportError:
|
||||
cp = None
|
||||
CUPY_AVAILABLE = False
|
||||
|
||||
# 缓存形态学操作的核，避免重复创建
|
||||
_morphology_kernels = {}
|
||||
|
||||
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)。
|
||||
@@ -19,30 +30,52 @@ def compute_foe(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample
|
||||
"""
|
||||
H, W = flow_field.shape[:2]
|
||||
|
||||
# 优化：使用视图而不是复制，减少内存分配
|
||||
u = flow_field[:, :, 0]
|
||||
v = flow_field[:, :, 1]
|
||||
|
||||
# 下采样并计算散度，然后插值回原尺寸
|
||||
if downsample_factor > 1:
|
||||
flow_ds = flow_field[::downsample_factor, ::downsample_factor, :]
|
||||
u_ds, v_ds = flow_ds[:, :, 0], flow_ds[:, :, 1]
|
||||
# 使用步进切片，避免复制
|
||||
u_ds = u[::downsample_factor, ::downsample_factor]
|
||||
v_ds = v[::downsample_factor, ::downsample_factor]
|
||||
|
||||
magnitude_ds = np.sqrt(u_ds**2 + v_ds**2)
|
||||
magnitude_ds = cv2.GaussianBlur(magnitude_ds, (0, 0), denoise_sigma)
|
||||
divergence_ds = np.gradient(u_ds, axis=1) + np.gradient(v_ds, axis=0)
|
||||
divergence_ds = cv2.GaussianBlur(divergence_ds, (0, 0), denoise_sigma)
|
||||
# 计算幅度（优化：使用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:
|
||||
u, v = flow_field[:, :, 0], flow_field[:, :, 1]
|
||||
# 计算幅度（优化：使用square减少中间数组）
|
||||
magnitude = np.sqrt(np.square(u) + np.square(v))
|
||||
if denoise_sigma > 0:
|
||||
magnitude = cv2.GaussianBlur(magnitude, (0, 0), denoise_sigma)
|
||||
|
||||
magnitude = np.sqrt(u**2 + v**2)
|
||||
magnitude = cv2.GaussianBlur(magnitude, (0, 0), denoise_sigma)
|
||||
divergence = np.gradient(u, axis=1) + np.gradient(v, axis=0)
|
||||
# 优化：合并梯度计算
|
||||
du_dx = np.gradient(u, axis=1)
|
||||
dv_dy = np.gradient(v, axis=0)
|
||||
divergence = du_dx + dv_dy
|
||||
|
||||
# 降噪：在计算散度之后对散度图进行平滑
|
||||
divergence = cv2.GaussianBlur(divergence, (0, 0), denoise_sigma)
|
||||
if denoise_sigma > 0:
|
||||
divergence = cv2.GaussianBlur(divergence, (0, 0), denoise_sigma)
|
||||
|
||||
divergence = np.where(divergence>0.02, divergence, 0)
|
||||
# 创建掩码
|
||||
# 优化：使用原地操作
|
||||
np.where(divergence > 0.02, divergence, 0, out=divergence)
|
||||
|
||||
# 创建掩码（使用原地操作优化）
|
||||
valid_mask = magnitude > min_magnitude
|
||||
expansion_mask = (valid_mask & (divergence > expansion_threshold))
|
||||
contraction_mask = (valid_mask & (divergence < -expansion_threshold))
|
||||
@@ -56,9 +89,153 @@ def compute_foe(flow_field, min_magnitude=1.0, expansion_threshold=1, downsample
|
||||
}
|
||||
|
||||
# 寻找扩张区域轮廓
|
||||
# 优化：缓存形态学操作的核
|
||||
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, np.ones(kernel_size, np.uint8))
|
||||
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['is_foe_found'] = True
|
||||
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:
|
||||
result['foe_point'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
|
||||
|
||||
return result
|
||||
|
||||
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 > 0.02, 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 = {
|
||||
'expansion_mask': expansion_mask,
|
||||
'contraction_mask': contraction_mask,
|
||||
'divergence_map': divergence,
|
||||
'foe_points': [0, 0],
|
||||
'is_foe_found': False
|
||||
}
|
||||
|
||||
# 寻找扩张区域轮廓（在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:
|
||||
|
||||
Reference in New Issue
Block a user