unknown
357 lines
16 KiB
Python
357 lines
16 KiB
Python
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 = {}
|
||
|
||
class Contact3D:
|
||
def __init__(self, config_contact=None, verbose=False):
|
||
|
||
# Load Config
|
||
self.config_contact = config_contact
|
||
self.flow_threshold = config_contact["flow_threshold"]
|
||
self.expansion_threshold = config_contact["expansion_threshold"]
|
||
self.downsample_factor = 1
|
||
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}")
|
||
|
||
# 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 _create_foe_result(self, 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(self, flow_field, return_all_contours=False):
|
||
"""
|
||
计算光流场的扩张焦点(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:
|
||
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
|