display without edge + depth map smoothing

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"):

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
+    Compute filtered density map from optical flow field.
-    # tmp = time.time()
+    
+    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])
+        zero_level = current_zero_level
-    # threshold = zero_level - (density.max() - zero_level) / 120
+    else:
-    threshold = zero_level
+        # 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
+    # Fill invalid areas using morphological closing to smooth transitions
-    density[np.where(invalid_ind > 0)] = zero_level  ## type 1 repair
+    # This fills small holes and provides better continuity
-    # density = guided_filter(density, density, 7, 0.005, s=4)
+    if np.any(invalid_ind > 0):
-    density = cv2.ximgproc.guidedFilter(density, density, 7, 0.005)
+        # 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_8U: "u1",   # unsigned 8-bit
-            cv2.CV_8S: "i1",
+            cv2.CV_8S: "i1",   # signed 8-bit
-            cv2.CV_16U: "u2",
+            cv2.CV_16U: "u2",  # unsigned 16-bit
-            cv2.CV_16S: "i2",
+            cv2.CV_16S: "i2",  # signed 16-bit
-            cv2.CV_32S: "i4",
+            cv2.CV_32S: "i4",  # signed 32-bit integer
-            cv2.CV_32F: "f4",
+            cv2.CV_32F: "f4",  # 32-bit float
-            cv2.CV_64F: "f8",
+            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 distance = dx*dx+dy*dy;  // Squared distance
-        float result = expf(-distance/twoSigma2);
+        float result = expf(-distance/twoSigma2);  // Gaussian exponential term
-        result = result/(twoSigma2*pi);
+        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;
+        // Get thread coordinates (each thread handles one source pixel)
-        int y = blockIdx.x;
+        int x = threadIdx.x;  // Column index
-        float dest_x = float(x)+flowX[y*x_size+x];
+        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<ceilf(dest_x+r_max); ++ind_x)
+        
-            for(int ind_y=ceilf(dest_y-r_max);ind_y<ceilf(dest_y+r_max); ++ind_y)
+        // Iterate over nearby grid points within r_max radius
+        for(int ind_x = ceilf(dest_x - r_max); ind_x < ceilf(dest_x + r_max); ++ind_x)
+            for(int ind_y = ceilf(dest_y - r_max); ind_y < ceilf(dest_y + r_max); ++ind_y)
             {
-                if(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;
+                // Compute distance from destination to grid point
-                atomicAdd(result+ind_y*x_size+ind_x, -gaussian(twoSigma2, dx,dy));
+                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),
+        flowx = cp.array(CudaArrayInterface(flowx), dtype=cp.float32, copy=True)
-                         dtype=cp.float32,
+        flowy = cp.array(CudaArrayInterface(flowy), dtype=cp.float32, copy=True)
-                         copy=True)
+    
-        flowy = cp.array(CudaArrayInterface(flowy),
+    # Initialize output array on GPU
-                         dtype=cp.float32,
-                         copy=True)
     result = cp.zeros((flowx.shape[0], flowx.shape[1]), dtype=cp.float32)
-    block_dim = (flowx.shape[1], )
+    
-    grid = (flowx.shape[0], )
+    # Configure CUDA kernel launch parameters
-    funcupy(grid, block_dim, (flowx, flowy, result, cp.int32(
+    block_dim = (flowx.shape[1], )  # Threads per block (one thread per column)
-        flowx.shape[1]), cp.int32(flowx.shape[0]), twoSigma2, r_max))
+    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()