diff --git a/example_contact.py b/example_contact.py index 8934dd1..163a5c4 100644 --- a/example_contact.py +++ b/example_contact.py @@ -3,7 +3,7 @@ os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" # Windows 下禁 import cv2 import orisys import numpy as np - +import time def main(): # 方式1: 从视频文件读取 @@ -46,7 +46,8 @@ def main(): cv2.imshow("magnitude map", mag_img) # 计算触碰点 - contact_result = sensor.get_contact(multiple_foe=False, min_magnitude=1, expansion_threshold=0.01,denoise_sigma=15) + contact_result = sensor.get_contact(min_magnitude=1, expansion_threshold=0.01,denoise_sigma=15) + # 在图像上触碰点和区域 image_with_foe = orisys.util.draw_foe(sensor.img.copy(), contact_result) cv2.imshow("Contact Centroid", image_with_foe) @@ -63,6 +64,8 @@ def main(): # 应用颜色映射（使用JET：蓝色=低散度，红色=高散度） cv2.imshow("Divergence Map", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET)) + + print(f"fps: {fps}") # ===================================================================== diff --git a/src/orisys/get_depth.py b/src/orisys/get_depth.py index 674a45d..568d195 100644 --- a/src/orisys/get_depth.py +++ b/src/orisys/get_depth.py @@ -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: diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index df9f830..b5e5ebe 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -249,7 +249,12 @@ class Sensor: if multiple_foe: return 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) else: - return compute_foe(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) + if self.cuda: + return 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) + + else: + return compute_foe(self.vnormal, min_magnitude=min_magnitude, expansion_threshold=expansion_threshold, downsample_factor=downsample_factor, kernel_size=kernel_size, denoise_sigma=denoise_sigma) + def get_depth(self, denoise_sigma=15): return getFilteredDensity(self.vnormal, use_cuda=self.cuda) @@ -327,13 +332,15 @@ class Sensor: self._capture_queue = Queue(maxsize=queue_size) self._flow_queue = Queue(maxsize=queue_size) self._result_queue = Queue(maxsize=queue_size) + self._contact_queue = Queue(maxsize=queue_size) self._stop_event = Event() self._async_enabled = True self._capture_thread = Thread(target=self._capture_worker, daemon=True) self._flow_thread = Thread(target=self._flow_worker, daemon=True) self._compute_thread = Thread(target=self._hhd_worker, daemon=True) - + self._contact_thread = Thread(target=self._contact_worker, daemon=True) + self._capture_thread.start() self._flow_thread.start() self._compute_thread.start() @@ -567,6 +574,18 @@ class Sensor: except Exception as e: if self.verbose: print(f"HHD error: {e}") + def _contact_worker(self): + """Stage 4: Contact computation.""" + while not self._stop_event.is_set(): + try: + img, flow = self._flow_queue.get(timeout=0.1) + contact_result = self.get_contact(flow) + self._queue_put(self._result_queue, contact_result) + except Empty: + continue + except Exception as e: + if self.verbose: print(f"Contact error: {e}") + # ========================================================================= # Private - Utilities # ========================================================================= diff --git a/test_foe_depth.py b/test_foe_depth.py index 33566f4..3d9041b 100644 --- a/test_foe_depth.py +++ b/test_foe_depth.py @@ -321,7 +321,7 @@ def main(): # 计算扩张焦点 # foe_result = orisys.util.compute_foe(vnormal,min_magnitude=1,expansion_threshold=0.01,downsample_factor=1) - foe_result = sensor.get_foe(multiple_foe=False, min_magnitude=1, expansion_threshold=0.01, downsample_factor=1,denoise_sigma=15) + foe_result = sensor.get_contact(multiple_foe=False, min_magnitude=1, expansion_threshold=0.01, downsample_factor=1,denoise_sigma=15) # 在图像上绘制焦点和区域 image_with_foe = orisys.util.draw_foe(sensor.img.copy(), foe_result) cv2.imshow("Focus of Expansion", image_with_foe)