diff --git a/config/compound_eye.png b/config/compound_eye.png index 880f357..2e4f334 100644 Binary files a/config/compound_eye.png and b/config/compound_eye.png differ diff --git a/example_contact.py b/example_contact.py index 0f7a8b2..8ec27ac 100644 --- a/example_contact.py +++ b/example_contact.py @@ -10,7 +10,7 @@ def main(): # sensor = orisys.Sensor("./s1.mp4", isstitch=True, cuda=True, verbose=True) # 方式2: 从摄像头读取（取消注释使用） - sensor = orisys.Sensor(0, isstitch=True, cuda=True, verbose=True) + sensor = orisys.Sensor(0, isstitch=True, cuda=False, verbose=True) print("\n按 'q' 退出程序，按 'r' 重置光流追踪器") @@ -37,7 +37,8 @@ def main(): cv2.imshow("形变场 Deformation", arrows) # 计算触碰点 - sensor.compute_contact(min_magnitude=0.1,expansion_threshold=0.01,denoise_sigma=15) + # sensor.compute_contact(min_magnitude=0.1,expansion_threshold=0.01,denoise_sigma=15) + sensor.compute_contact() # 在图像上触碰点和区域 image_with_foe = orisys.util.draw_contact(sensor.img.copy(), sensor.contour, sensor.centroid) diff --git a/src/orisys/Stitcher.py b/src/orisys/Stitcher.py index c70f500..b6f0441 100644 --- a/src/orisys/Stitcher.py +++ b/src/orisys/Stitcher.py @@ -6,15 +6,15 @@ import cv2 class Stitcher: - def __init__(self,isstitch=False, n=[3,3], index=np.array([1,2,3,6,7,8,11,12,13])): + def __init__(self,isstitch=False, n=[3,3], index=np.array([1,2,3,6,7,8,11,12,13]), config_stitcher=None): # Settings self.isstitch = isstitch - + self.config_stitcher = config_stitcher # Find Boxes Related - self.min_area_frac = 0.01 - self.max_area_frac = 0.05 - self.max_aspect_ratio = 100 + self.min_area_frac = config_stitcher["min_area_frac"] + self.max_area_frac = config_stitcher["max_area_frac"] + self.max_aspect_ratio = config_stitcher["max_aspect_ratio"] # stitch related self.stitch_valid = False diff --git a/src/orisys/configs/prototype.json b/src/orisys/configs/prototype.json index f2f81c8..187a8ea 100644 --- a/src/orisys/configs/prototype.json +++ b/src/orisys/configs/prototype.json @@ -17,7 +17,8 @@ }, "contact": { "min_contour_area": 3000, - "min_magnitude": 0.02, + "flow_threshold":1, + "div_threshold":0.02, "expansion_threshold": 0.01, "downsample_factor": 1, "kernel_size": [5,5], diff --git a/src/orisys/get_depth.py b/src/orisys/get_depth.py index 02d9503..dfdd0d6 100644 --- a/src/orisys/get_depth.py +++ b/src/orisys/get_depth.py @@ -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 diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index 7d69009..5c7f992 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -20,7 +20,7 @@ from .tracker import Tracker from .nHHD import nHHD # from .getDensity import getFilteredDensity from . import configs -from .get_depth import * +from .get_depth import Contact3D class Sensor: @@ -132,8 +132,9 @@ class Sensor: # Initialize components - self._stitcher = Stitcher(isstitch=isstitch, index=self.config_index) + self._stitcher = Stitcher(isstitch=isstitch, index=self.config_index, config_stitcher=self.config_stitcher) self.tracker = Tracker(adaptive=True, cuda=cuda) + self.contact3d = Contact3D(config_contact=self.config_contact, verbose=self.verbose) self.hhd = None # Setup video capture @@ -261,14 +262,15 @@ class Sensor: def compute_contact(self, multiple_foe=False, min_contour_area=3000, min_magnitude=1.0, expansion_threshold=0.01, downsample_factor=1, kernel_size=(5,5), denoise_sigma=15): if multiple_foe: - result = compute_multiple_foe(self.vnormal, min_contour_area=min_contour_area, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) - + # result = compute_multiple_foe(self.vnormal, min_contour_area=min_contour_area, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) + result = self.contact3d.compute_multiple_foe(self.vnormal) else: if self.cuda: - result = compute_foe_gpu(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) - + # result = compute_foe_gpu(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) + result = self.contact3d.compute_foe_gpu(self.vnormal) else: - result = compute_foe(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) + # result = compute_foe(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) + result = self.contact3d.compute_foe(self.vnormal) self.centroid = result['centroid'] self.centroid_found = result['centroid_found'] @@ -457,9 +459,8 @@ class Sensor: self.config_index = cam["index"] self.config_pos_default = cam["pos_default"] - config_stitcher = self.config["stitcher"] - - config_contact = self.config["contact"] + self.config_stitcher = self.config["stitcher"] + self.config_contact = self.config["contact"]