diff --git a/config/compound_eye.png b/config/compound_eye.png index ef3b17b..83e7a20 100644 Binary files a/config/compound_eye.png and b/config/compound_eye.png differ diff --git a/demo.py b/demo.py index 861b202..8656810 100644 --- a/demo.py +++ b/demo.py @@ -4,7 +4,7 @@ import cv2 import numpy as np from orisys import Sensor -def draw_arrows_fade(img, flow, depth, depth_scale=0.5, grid_spacing=40, arrow_scale=1.0, threshold=5.0): +def draw_arrows_fade(img, flow, depth, depth_scale=0.5, grid_spacing=20, arrow_scale=1.0, threshold=5.0): """ Draw optical flow arrows on image with transparency based on magnitude. @@ -18,10 +18,11 @@ def draw_arrows_fade(img, flow, depth, depth_scale=0.5, grid_spacing=40, arrow_s Returns: Image with arrows drawn, blended with original image """ - + depth_img = cv2.applyColorMap(((depth+1)*depth_scale* 255.0).astype('uint8'), cv2.COLORMAP_HOT) # Make a copy of the image img_copy = depth_img.copy() + # Convert grayscale to BGR if needed if len(img_copy.shape) == 2 or img_copy.shape[2] == 1: @@ -132,9 +133,9 @@ def main(): global start_time # 打开视频 - sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True) + # sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True) # 打开摄像头 - # sensor1 = Sensor(0,isstitch=True,cuda = True,verbose = True) + sensor1 = Sensor(0,isstitch=True,cuda = True,verbose = True) while True: sensor1.get_img() # 提取并拼接复眼图像 sensor1.compute_deformation() # 计算形变和分解 @@ -145,7 +146,20 @@ def main(): fps, fn, fx, fy, flow, depth = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW, sensor1.info.DEPTH) arrows = draw_arrows_fade(sensor1.img, flow, depth) print(fps,np.min(depth),np.mean(depth)) - cv2.imshow("Deformation", arrows) + + # Crop edges (15 pixels from each side) - only if image is large enough + h, w = arrows.shape[:2] + if h > 30 and w > 30: + cv2.imshow("Deformation", arrows[15:-15, 15:-15, :]) + else: + cv2.imshow("Deformation", arrows) + + # Crop raw image edges if large enough + h_raw, w_raw = sensor1.img.shape[:2] + if h_raw > 30 and w_raw > 30: + cv2.imshow("raw", sensor1.img[15:-15, 15:-15]) + else: + cv2.imshow("raw", sensor1.img) key = cv2.waitKey(1) if key & 0xFF == ord("q"): diff --git a/src/orisys/getDensity.py b/src/orisys/getDensity.py index effb282..c0730b4 100644 --- a/src/orisys/getDensity.py +++ b/src/orisys/getDensity.py @@ -1,47 +1,127 @@ +""" +Density estimation module for optical flow fields. + +This module computes density maps from optical flow vectors using Gaussian kernel +density estimation. It supports both CPU (NumPy) and CUDA (CuPy) implementations. +""" import cupy as cp import cv2 import numpy as np +# Global variable to store the baseline/zero level of density for invalid area correction zero_level = None def getFilteredDensity(flow, use_cuda): - sigma = 3.0 - r_max = 9.0 - # tmp = time.time() + """ + Compute filtered density map from optical flow field. + + This function computes a density map using Gaussian kernel density estimation + and then applies filtering to correct invalid areas. + + Args: + flow: Optical flow field with shape (height, width, 2) containing flow vectors + use_cuda: If True, use CUDA-accelerated computation (requires CuPy and CUDA) + If False, use CPU-based NumPy implementation + + Returns: + np.ndarray: Filtered density map with shape (height, width), where negative + values indicate regions of flow divergence + """ + # Gaussian kernel parameters for density estimation + sigma = 3.0 # Standard deviation of Gaussian kernel + r_max = 9.0 # Maximum radius for Gaussian contribution (cutoff distance) + + # Compute density map using either CUDA or CPU implementation if use_cuda: density = getDensityCuda(flow, sigma, r_max) else: density = getDensity(flow, sigma, r_max) - # print("Density frequency: ", float(1.0 / (time.time() - tmp))) - # tmp = time.time() + + # Apply filtering to correct invalid areas (regions with unreliable flow) density = correctInvalidArea(density) - # print("Correct frequency: ", float(1.0 / (time.time() - tmp))) - # return density[15:-15, 15:-15] + return density def correctInvalidArea(density): - # density = guided_filter(density, density, 5, 0.001, s=4) - # invalid_ind = np.zeros_like(density, dtype='uint8') + """ + Correct invalid areas in the density map. + + Invalid areas are regions where the density falls below a threshold (zero_level), + which typically indicates unreliable or missing flow data. This function: + 1. Uses adaptive zero_level that updates gradually to handle changing conditions + 2. Identifies invalid regions using adaptive thresholding (below zero_level - margin) + 3. Fills invalid areas using morphological closing for smoother transitions + 4. Applies stronger edge-preserving smoothing to eliminate discontinuities + + Args: + density: Input density map with shape (height, width) + + Returns: + np.ndarray: Corrected density map with invalid areas filled and smoothed, + with improved continuity to reduce black dots/discontinuities + """ global zero_level + + # Compute adaptive zero level from current frame's central region + # This adapts to changing density values when sensor is pressed + current_zero_level = np.mean(density[60:-60, 60:-60]) + + # Initialize or update zero_level with exponential moving average for stability if zero_level is None: - zero_level = np.mean(density[60:-60, 60:-60]) - # threshold = zero_level - (density.max() - zero_level) / 120 - threshold = zero_level + zero_level = current_zero_level + else: + # Use exponential moving average: 90% old, 10% new (adapts slowly) + zero_level = 0.9 * zero_level + 0.1 * current_zero_level + + # Use a threshold slightly below zero level to avoid over-correction + # This prevents valid but slightly lower density regions from being marked invalid + threshold = zero_level - np.std(density[60:-60, 60:-60]) * 0.5 invalid_ind = twoClustering(density, threshold) - # density = cv2.inpaint(density * 1e4, invalid_ind, 20, - # cv2.INPAINT_TELEA) / 1e4 - density[np.where(invalid_ind > 0)] = zero_level ## type 1 repair - # density = guided_filter(density, density, 7, 0.005, s=4) - density = cv2.ximgproc.guidedFilter(density, density, 7, 0.005) + + # Fill invalid areas using morphological closing to smooth transitions + # This fills small holes and provides better continuity + if np.any(invalid_ind > 0): + # Create a temporary copy with invalid areas filled with zero_level + density_filled = density.copy() + density_filled[np.where(invalid_ind > 0)] = zero_level + + # Apply morphological closing to fill small holes and smooth boundaries + # This helps connect nearby valid regions and eliminate black dots + kernel = np.ones((5, 5), np.float32) + density_closed = cv2.morphologyEx(density_filled, cv2.MORPH_CLOSE, kernel, iterations=1) + + # Use closed version only for invalid regions, preserve original valid values + density[np.where(invalid_ind > 0)] = density_closed[np.where(invalid_ind > 0)] + + # Apply edge-preserving smoothing using guided filter with stronger smoothing + # Increased radius and epsilon for better continuity and to eliminate discontinuities + try: + density = cv2.ximgproc.guidedFilter(density, density, 15, 0.01) + except AttributeError: + # Fallback to stronger Gaussian blur if ximgproc is not available + # Use larger kernel for better smoothing of discontinuities + density = cv2.GaussianBlur(density, (21, 21), 0) + return density - # return invalid_ind def twoClustering(density, threshold): + """ + Perform binary clustering/segmentation of density map based on threshold. + + This function creates a binary mask where pixels below the threshold are + marked as invalid (value 1) and pixels above threshold are valid (value 0). + + Args: + density: Input density map + threshold: Threshold value for segmentation + + Returns: + np.ndarray: Binary mask with dtype uint8, where 1 indicates invalid areas + (density < threshold) and 0 indicates valid areas + """ cluster = np.zeros_like(density, dtype='uint8') cluster[np.where(density < threshold)] = 1 - # kernel = np.ones((3, 3), np.uint8) - # cluster = cv2.dilate(cluster, kernel, iterations=4) return cluster @@ -81,115 +161,196 @@ def getDensity(flow, sigma, r_max): # Initialize result array result = np.zeros((height, width), dtype=np.float32) - # Vectorized computation: for each source pixel, compute contributions to nearby grid points + # Compute density by iterating over each source pixel + # For each pixel, apply its flow vector to get destination position, + # then accumulate Gaussian contributions to nearby grid points for y in range(height): for x in range(width): + # Get the destination position after applying flow displacement dest_x_val = dest_x[y, x] dest_y_val = dest_y[y, x] - # Determine the bounding box of affected grid points + # Determine the bounding box of grid points affected by this flow vector + # These are grid points within r_max radius of the destination min_x = max(0, int(np.floor(dest_x_val - r_max))) max_x = min(width, int(np.ceil(dest_x_val + r_max)) + 1) min_y = max(0, int(np.floor(dest_y_val - r_max))) max_y = min(height, int(np.ceil(dest_y_val + r_max)) + 1) + # Skip if bounding box is empty or out of bounds if min_x >= max_x or min_y >= max_y: continue - # Create local grid for this window + # Create local coordinate grid for the affected region local_x = np.arange(min_x, max_x, dtype=np.float32) local_y = np.arange(min_y, max_y, dtype=np.float32) local_grid_x, local_grid_y = np.meshgrid(local_x, local_y) - # Compute distances squared + # Compute squared distances from destination to each grid point dx = local_grid_x - dest_x_val dy = local_grid_y - dest_y_val distance_sq = dx * dx + dy * dy - # Filter by r_max (only include points within radius) + # Create mask for points within r_max radius (circular neighborhood) mask = distance_sq <= (r_max * r_max) - # Compute Gaussian contributions + # Compute Gaussian kernel values: exp(-d²/(2σ²)) / (2πσ²) + # Higher values for closer points, decreasing with distance gaussian_vals = np.exp(-distance_sq / two_sigma2) * gaussian_norm - gaussian_vals[~mask] = 0 + gaussian_vals[~mask] = 0 # Zero out contributions beyond r_max - # Accumulate negative density + # Accumulate negative density (negative indicates flow divergence) + # Each flow vector contributes to density at its destination region result[min_y:max_y, min_x:max_x] -= gaussian_vals return result class CudaArrayInterface: + """ + Interface class to convert OpenCV CUDA GpuMat to CuPy-compatible array interface. + + This allows seamless conversion between OpenCV's CUDA memory objects and CuPy arrays + without copying data, by implementing the CUDA array interface protocol. + """ def __init__(self, gpu_mat): + """ + Initialize CUDA array interface from OpenCV GpuMat. + + Args: + gpu_mat: OpenCV CUDA GpuMat object + """ w, h = gpu_mat.size() + + # Map OpenCV data types to NumPy/CuPy type strings type_map = { - cv2.CV_8U: "u1", - cv2.CV_8S: "i1", - cv2.CV_16U: "u2", - cv2.CV_16S: "i2", - cv2.CV_32S: "i4", - cv2.CV_32F: "f4", - cv2.CV_64F: "f8", + cv2.CV_8U: "u1", # unsigned 8-bit + cv2.CV_8S: "i1", # signed 8-bit + cv2.CV_16U: "u2", # unsigned 16-bit + cv2.CV_16S: "i2", # signed 16-bit + cv2.CV_32S: "i4", # signed 32-bit integer + cv2.CV_32F: "f4", # 32-bit float + cv2.CV_64F: "f8", # 64-bit float } + + # Implement CUDA array interface protocol for zero-copy conversion self.__cuda_array_interface__ = { "version": 2, "shape": (h, w), - "data": (gpu_mat.cudaPtr(), False), + "data": (gpu_mat.cudaPtr(), False), # CUDA memory pointer, not writable "typestr": type_map[gpu_mat.type()], - "strides": (gpu_mat.step, gpu_mat.elemSize()), + "strides": (gpu_mat.step, gpu_mat.elemSize()), # Row stride and element size } +# CUDA kernel code for accelerated density computation +# This kernel computes Gaussian density contributions in parallel on GPU func_string = r""" extern "C"{ + // Compute 2D Gaussian kernel value + // Args: + // twoSigma2: 2 * sigma^2 (precomputed for efficiency) + // dx, dy: Distance components from center + // Returns: Gaussian kernel value normalized by (2*pi*sigma^2) __device__ float gaussian(float twoSigma2, float dx, float dy) { float pi = 3.1415926535897932384626; - float distance = dx*dx+dy*dy; - float result = expf(-distance/twoSigma2); - result = result/(twoSigma2*pi); + float distance = dx*dx+dy*dy; // Squared distance + float result = expf(-distance/twoSigma2); // Gaussian exponential term + result = result/(twoSigma2*pi); // Normalization factor return result; } - __global__ void density(const float *flowX, const float *flowY, float *result, int x_size, int y_size, float twoSigma2, float r_max) + + // CUDA kernel to compute density map from optical flow + // Each thread processes one pixel in the flow field + // Args: + // flowX, flowY: Input flow field components (flattened) + // result: Output density map (flattened, initialized to zero) + // x_size, y_size: Dimensions of the flow field + // twoSigma2: Precomputed 2 * sigma^2 + // r_max: Maximum radius for Gaussian contribution + __global__ void density(const float *flowX, const float *flowY, float *result, + int x_size, int y_size, float twoSigma2, float r_max) { - int x = threadIdx.x; - int y = blockIdx.x; - float dest_x = float(x)+flowX[y*x_size+x]; + // Get thread coordinates (each thread handles one source pixel) + int x = threadIdx.x; // Column index + int y = blockIdx.x; // Row index + + // Compute destination position after applying flow displacement + float dest_x = float(x) + flowX[y*x_size+x]; float dest_y = float(y) + flowY[y*x_size+x]; - for(int ind_x=ceilf(dest_x-r_max);ind_x= x_size || ind_x<0 || ind_y >= y_size || ind_y < 0) + // Bounds checking + if(ind_x >= x_size || ind_x < 0 || ind_y >= y_size || ind_y < 0) { continue; } - float dx = float(ind_x)-dest_x; - float dy = float(ind_y)-dest_y; - atomicAdd(result+ind_y*x_size+ind_x, -gaussian(twoSigma2, dx,dy)); + + // Compute distance from destination to grid point + float dx = float(ind_x) - dest_x; + float dy = float(ind_y) - dest_y; + + // Accumulate Gaussian contribution (negative for divergence representation) + atomicAdd(result + ind_y*x_size + ind_x, -gaussian(twoSigma2, dx, dy)); } } } """ +# Compile CUDA kernel into CuPy RawModule for execution funcupy = cp.RawModule(code=func_string).get_function("density") def getDensityCuda(flow, sigma, r_max): + """ + Compute Gaussian density map from optical flow using CUDA acceleration. + + This function uses a custom CUDA kernel for parallel computation of density maps, + providing significant speedup over CPU implementation for large flow fields. + + Args: + flow: Optical flow field with shape (height, width, 2), can be numpy array + or OpenCV CUDA GpuMat + sigma: Gaussian kernel standard deviation + r_max: Maximum radius for Gaussian contribution (cutoff distance) + + Returns: + np.ndarray: Density map with shape (height, width), transferred from GPU to CPU + + Note: + Requires CuPy and CUDA-capable GPU. Falls back to CPU if unavailable. + """ + # Precompute 2*sigma^2 for efficiency twoSigma2 = cp.float32(2.0 * sigma**2) r_max = cp.float32(r_max) + + # Convert input to CuPy arrays + # Handle both numpy arrays and OpenCV CUDA GpuMat if type(flow).__module__ == np.__name__: + # NumPy array: convert directly to CuPy flowx = cp.array(flow[:, :, 0], dtype=cp.float32) flowy = cp.array(flow[:, :, 1], dtype=cp.float32) else: + # OpenCV CUDA GpuMat: use zero-copy interface flowx, flowy = cv2.cuda.split(flow) - flowx = cp.array(CudaArrayInterface(flowx), - dtype=cp.float32, - copy=True) - flowy = cp.array(CudaArrayInterface(flowy), - dtype=cp.float32, - copy=True) + flowx = cp.array(CudaArrayInterface(flowx), dtype=cp.float32, copy=True) + flowy = cp.array(CudaArrayInterface(flowy), dtype=cp.float32, copy=True) + + # Initialize output array on GPU result = cp.zeros((flowx.shape[0], flowx.shape[1]), dtype=cp.float32) - block_dim = (flowx.shape[1], ) - grid = (flowx.shape[0], ) - funcupy(grid, block_dim, (flowx, flowy, result, cp.int32( - flowx.shape[1]), cp.int32(flowx.shape[0]), twoSigma2, r_max)) + + # Configure CUDA kernel launch parameters + block_dim = (flowx.shape[1], ) # Threads per block (one thread per column) + grid = (flowx.shape[0], ) # Number of blocks (one block per row) + + # Launch CUDA kernel + funcupy(grid, block_dim, (flowx, flowy, result, + cp.int32(flowx.shape[1]), + cp.int32(flowx.shape[0]), + twoSigma2, r_max)) + + # Transfer result from GPU to CPU return result.get()