diff --git a/config/compound_eye.png b/config/compound_eye.png index 3f6510b..ef3b17b 100644 Binary files a/config/compound_eye.png and b/config/compound_eye.png differ diff --git a/example1_basic_pipeline.py b/example1_basic_pipeline.py index 21d291a..1c956d0 100644 --- a/example1_basic_pipeline.py +++ b/example1_basic_pipeline.py @@ -19,11 +19,11 @@ def main(): # 步骤1：创建传感器对象 (必要步骤)，config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径） # 方式1: 从视频文件读取 - # sensor = orisys.Sensor("s1.mp4", config_name="prototype", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) + sensor = orisys.Sensor("s1.mp4", config_name="prototype", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) # 方式2: 从摄像头读取（取消注释使用） - sensor = orisys.Sensor(0, config_name="./config/lens_test.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) + # sensor = orisys.Sensor(0, config_name="./config/lens_test.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) print("\n按 'q' 退出程序，按 'r' 重置光流追踪器") @@ -33,12 +33,12 @@ def main(): # 步骤 2.1: 获取并拼接图像 (必要步骤) sensor.get_img() print("img_size_raw:", sensor.img_size_raw) - if sensor.detect_motion(threshold=2): - print("motion detected") - # 步骤 2.2: 计算形变场 (必要步骤) - sensor.compute_deformation() - # 步骤 2.3: 计算触碰点 - sensor.compute_contact() + # if sensor.detect_motion(threshold=2): + # print("motion detected") + # 步骤 2.2: 计算形变场 (必要步骤) + sensor.compute_deformation() + # 步骤 2.3: 计算触碰点 + sensor.compute_contact() #步骤 2.4: 读取信息 diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py index 50a3e7e..fb02570 100644 --- a/src/orisys/HHD.py +++ b/src/orisys/HHD.py @@ -102,7 +102,11 @@ class nHHD(object): if (self.dim != 2) and (self.dim != 3): raise ValueError("nHHD Solver works only for 2 and 3 dimensions") - self.psolver.prepare(False) + self.psolver.prepare(False) + + # Example pattern for stream usage: + self.stream1 = cp.cuda.Stream() + self.stream2 = cp.cuda.Stream() # create the 3-component decomposition def decompose(self, vfield, verbose=0): @@ -228,20 +232,21 @@ class nHHD(object): dudy, dudx = cp.gradient(vfield_gpu[:,:,0], dx0, dx1) dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], dx0, dx1) - div_gpu = dudx + dvdy - curlw_gpu = dvdx - dudy + div_gpu = cp.add(dudx, dvdy) + curlw_gpu = cp.subtract(dvdx, dudy) if verbose > 0: print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max())) print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max())) + # Compute potentials on GPU (solve_gpu returns numpy, so we convert back) - nD_cpu = self.psolver.solve_gpu(cp.asnumpy(div_gpu), verbose > 0) - nRu_cpu = self.psolver.solve_gpu(cp.asnumpy(curlw_gpu), verbose > 0) - - # Keep potentials on GPU - nD_gpu = cp.asarray(nD_cpu, dtype=cp.float32) - nRu_gpu = cp.asarray(nRu_cpu, dtype=cp.float32) + with self.stream1: + nD_gpu = self.psolver._solve_gpu_internal(div_gpu) + with self.stream2: + nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu) + self.stream1.synchronize() + self.stream2.synchronize() if verbose > 0: print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max())) @@ -256,15 +261,15 @@ class nHHD(object): # Compute rotated gradient on GPU rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1) - rdy *= -1.0 + rdy = cp.multiply(rdy, -1.0) r_gpu = cp.stack((rdy, rdx), axis=-1) if verbose > 0: print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max())) # Compute h on GPU - h_gpu = d_gpu + r_gpu - h_gpu = vfield_gpu - h_gpu + h_gpu = cp.add(d_gpu, r_gpu) + h_gpu = cp.subtract(vfield_gpu, h_gpu) if verbose > 0: print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max())) @@ -276,11 +281,11 @@ class nHHD(object): self._nRu_gpu = nRu_gpu # Also store CPU copies for compatibility - self.d = cp.asnumpy(d_gpu) - self.h = cp.asnumpy(h_gpu) - self.r = cp.asnumpy(r_gpu) - self.nRu = cp.asnumpy(nRu_gpu) - self.nD = nD_cpu + # self.d = cp.asnumpy(d_gpu) + # self.h = cp.asnumpy(h_gpu) + # self.r = cp.asnumpy(r_gpu) + # self.nRu = cp.asnumpy(nRu_gpu) + # self.nD = cp.asnumpy(nD_gpu) return (d_gpu, h_gpu, r_gpu, nRu_gpu) @@ -375,13 +380,18 @@ class nHHD(object): return self.get_force(vfield, verbose) # Decompose the vector field - self.decompose(vfield, verbose) + self.decompose_gpu(vfield, verbose) # Keep data on GPU - avoid transfers by computing directly - d_gpu = cp.asarray(self.d, dtype=cp.float32) - h_gpu = cp.asarray(self.h, dtype=cp.float32) - r_gpu = cp.asarray(self.r, dtype=cp.float32) - nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32) + # d_gpu = cp.asarray(self.d, dtype=cp.float32) + # h_gpu = cp.asarray(self.h, dtype=cp.float32) + # r_gpu = cp.asarray(self.r, dtype=cp.float32) + # nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32) + + d_gpu = self._d_gpu + h_gpu = self._h_gpu + r_gpu = self._r_gpu + nRu_gpu = self._nRu_gpu # Use reduction kernels for better performance # Compute Sn using L1 norm (sum of absolute values) @@ -410,10 +420,17 @@ class nHHD(object): r_y = r_gpu[:, :, 1] # Compute torque with fused sum - Stor = float(cp.sum( - (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + - (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x - )) + # Stor = float(cp.sum( + # (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + + # (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x + # )) + # Replace lines 413-416: + if KERNEL_AVAILABLE and _torque_kernel is not None: + torque = _torque_kernel( + r_x, r_y, mgrid_0, mgrid_1, + float(max_0), float(max_1), float(min_0), float(min_1) + ) + Stor = float(cp.sum(torque)) self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32) @@ -430,12 +447,14 @@ class nHHD(object): cp.get_default_memory_pool().free_all_blocks() def get_shear(self): + if not CUPY_AVAILABLE: + return self.h + return cp.asnumpy(self._h_gpu) - return self.h - def get_normal(self): - - return self.d + if not CUPY_AVAILABLE: + return self.d + return cp.asnumpy(self._d_gpu) # --------------------------------------------------------------------------- @@ -574,72 +593,72 @@ class StructuredGrid(object): return (dudx, dwdy, dudz, dvdx) - def divcurl_gpu(self, vfield, verbose=False): - """ - GPU-accelerated computation of divergence and curl of vector field. - Uses CuPy for faster gradient computation on GPU. - """ - if not CUPY_AVAILABLE: - if verbose: - print('CuPy not available, falling back to CPU') - return self.divcurl(vfield, verbose) + # def divcurl_gpu(self, vfield, verbose=False): + # """ + # GPU-accelerated computation of divergence and curl of vector field. + # Uses CuPy for faster gradient computation on GPU. + # """ + # if not CUPY_AVAILABLE: + # if verbose: + # print('CuPy not available, falling back to CPU') + # return self.divcurl(vfield, verbose) - if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims): - raise ValueError("Dimensions of vector field should match that of the grid") + # if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims): + # raise ValueError("Dimensions of vector field should match that of the grid") - if verbose: - print(' Computing div and curl on GPU...', end ='') - sys.stdout.flush() - mtimer = Timer() + # if verbose: + # print(' Computing div and curl on GPU...', end ='') + # sys.stdout.flush() + # mtimer = Timer() - if self.dim == 2: - # Transfer to GPU - vfield_gpu = cp.asarray(vfield, dtype=cp.float32) + # if self.dim == 2: + # # Transfer to GPU + # vfield_gpu = cp.asarray(vfield, dtype=cp.float32) - # Compute gradients on GPU - dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1]) - dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1]) + # # Compute gradients on GPU + # dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1]) + # dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1]) - # Compute div and curl on GPU - div_gpu = dudx + dvdy - curl_gpu = dvdx - dudy + # # Compute div and curl on GPU + # div_gpu = dudx + dvdy + # curl_gpu = dvdx - dudy - # Transfer back to CPU - div = cp.asnumpy(div_gpu) - curl = cp.asnumpy(curl_gpu) + # # Transfer back to CPU + # div = cp.asnumpy(div_gpu) + # curl = cp.asnumpy(curl_gpu) - if verbose: - print(' Done!', end ='') - mtimer.end() + # if verbose: + # print(' Done!', end ='') + # mtimer.end() - return (div, curl) + # return (div, curl) - elif self.dim == 3: - # Transfer to GPU - vfield_gpu = cp.asarray(vfield, dtype=cp.float32) + # elif self.dim == 3: + # # Transfer to GPU + # vfield_gpu = cp.asarray(vfield, dtype=cp.float32) - # Compute gradients on GPU - dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2]) - dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2]) - dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2]) + # # Compute gradients on GPU + # dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2]) + # dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2]) + # dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2]) - # Compute div and curl components on GPU - div_gpu = dudx + dvdy + dwdz - curlu_gpu = dwdy - dvdz - curlv_gpu = dudz - dwdx - curlw_gpu = dvdx - dudy + # # Compute div and curl components on GPU + # div_gpu = dudx + dvdy + dwdz + # curlu_gpu = dwdy - dvdz + # curlv_gpu = dudz - dwdx + # curlw_gpu = dvdx - dudy - # Transfer back to CPU - div = cp.asnumpy(div_gpu) - curlu = cp.asnumpy(curlu_gpu) - curlv = cp.asnumpy(curlv_gpu) - curlw = cp.asnumpy(curlw_gpu) + # # Transfer back to CPU + # div = cp.asnumpy(div_gpu) + # curlu = cp.asnumpy(curlu_gpu) + # curlv = cp.asnumpy(curlv_gpu) + # curlw = cp.asnumpy(curlw_gpu) - if verbose: - print(' Done!', end ='') - mtimer.end() + # if verbose: + # print(' Done!', end ='') + # mtimer.end() - return (div, curlu, curlv, curlw) + # return (div, curlu, curlv, curlw) def curl3D(self, vfield, verbose=False): @@ -1187,7 +1206,7 @@ class PoissonSolver(object): # -------------------------------------------------------------------------- # GPU-accelerated solve using CuPy # -------------------------------------------------------------------------- - def solve_gpu(self, f, verbose=False): + def solve_gpu(self, f, verbose=False, return_gpu=False): """ GPU-accelerated Poisson solver using CuPy. Falls back to CPU if CuPy is not available. @@ -1240,111 +1259,151 @@ class PoissonSolver(object): # Scale by grid spacing (on GPU) p_gpu *= self._gpu_cache['gdx_prod'] + + if verbose: + print('Poisson solver (GPU) finished') + gtimer.end() # Transfer result back to CPU - p = cp.asnumpy(p_gpu) + if return_gpu: + return p_gpu # Return CuPy array + return cp.asnumpy(p_gpu) + # p = cp.asnumpy(p_gpu) elif self.stype == 'S': raise SyntaxError('Spatial solver not suggested for structured grids') # ---------------------------------------------------------------------- # unstructured points - not optimized for GPU - elif self.ptype == 'P': - if verbose: - print(' - GPU solver not optimized for unstructured points, using CPU') - return self.solve(f, verbose) + # elif self.ptype == 'P': + # if verbose: + # print(' - GPU solver not optimized for unstructured points, using CPU') + # return self.solve(f, verbose) # ---------------------------------------------------------------------- - if verbose: - print('Poisson solver (GPU) finished') - gtimer.end() - - return p + + # return p - def solve_gpu_batch(self, f_list, verbose=False): + def _solve_gpu_internal(self, f_gpu): + """Internal GPU solve that accepts and returns GPU arrays.""" """ - Batch GPU-accelerated Poisson solver for multiple inputs. - Solves multiple Poisson equations in a single GPU call. + GPU-accelerated Poisson solver using CuPy. + Falls back to CPU if CuPy is not available. Args: - f_list: List of input fields (numpy arrays) - verbose: Print timing information + f_gpu: Input field (numpy array) Returns: - List of solutions (numpy arrays on CPU) - """ - if not CUPY_AVAILABLE: - if verbose: - print("CuPy not available, falling back to CPU solver") - return [self.solve(f, verbose) for f in f_list] - - if not self.ready: - self.prepare(verbose) - - gtimer = Timer() - if verbose: - print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)') - - # Verify all inputs have same shape - fshape = f_list[0].shape - for f in f_list: - if f.shape != fshape: - raise ValueError("All input fields must have the same shape for batch processing") + p_gpu: Solution (CuPy array) + """ # ---------------------------------------------------------------------- # regular grid if self.ptype == 'G': - if self.gdims != fshape: - raise ValueError("Shape of function should match shape of grid") # convolution in frequency domain using GPU if self.stype == 'F': # Initialize GPU cache if needed if self._gpu_cache is None: - if verbose: - print(' - Transferring Green\'s function to GPU...') self._gpu_cache = { 'G_gpu': cp.asarray(self.G, dtype=cp.float32), 'gdx_prod': numpy.prod(self.gdx) } - results = [] + # Perform FFT convolution on GPU + # CuPy's fftconvolve is GPU-accelerated + p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same') - # Process all inputs in batch - for i, f in enumerate(f_list): - # Transfer input to GPU - f_gpu = cp.asarray(f, dtype=cp.float32) - - # Perform FFT convolution on GPU - p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same') - - # Scale by grid spacing (on GPU) - p_gpu *= self._gpu_cache['gdx_prod'] - - # Transfer result back to CPU - results.append(cp.asnumpy(p_gpu)) - - if verbose: - print(f' - Solved {i+1}/{len(f_list)}') - - if verbose: - print('Poisson solver (GPU batch) finished') - gtimer.end() - - return results - - elif self.stype == 'S': - raise SyntaxError('Spatial solver not suggested for structured grids') + # Scale by grid spacing (on GPU) + p_gpu *= self._gpu_cache['gdx_prod'] + + return p_gpu # Return CuPy array + + + # def solve_gpu_batch(self, f_list, verbose=False): + # """ + # Batch GPU-accelerated Poisson solver for multiple inputs. + # Solves multiple Poisson equations in a single GPU call. - # ---------------------------------------------------------------------- - # unstructured points - not optimized for GPU - elif self.ptype == 'P': - if verbose: - print(' - GPU solver not optimized for unstructured points, using CPU') - return [self.solve(f, verbose) for f in f_list] + # Args: + # f_list: List of input fields (numpy arrays) + # verbose: Print timing information + + # Returns: + # List of solutions (numpy arrays on CPU) + # """ + # if not CUPY_AVAILABLE: + # if verbose: + # print("CuPy not available, falling back to CPU solver") + # return [self.solve(f, verbose) for f in f_list] - # If we reach here, it's an unsupported case - raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver") + # if not self.ready: + # self.prepare(verbose) + + # gtimer = Timer() + # if verbose: + # print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)') + + # # Verify all inputs have same shape + # fshape = f_list[0].shape + # for f in f_list: + # if f.shape != fshape: + # raise ValueError("All input fields must have the same shape for batch processing") + + # # ---------------------------------------------------------------------- + # # regular grid + # if self.ptype == 'G': + # if self.gdims != fshape: + # raise ValueError("Shape of function should match shape of grid") + + # # convolution in frequency domain using GPU + # if self.stype == 'F': + # # Initialize GPU cache if needed + # if self._gpu_cache is None: + # if verbose: + # print(' - Transferring Green\'s function to GPU...') + # self._gpu_cache = { + # 'G_gpu': cp.asarray(self.G, dtype=cp.float32), + # 'gdx_prod': numpy.prod(self.gdx) + # } + + # results = [] + + # # Process all inputs in batch + # for i, f in enumerate(f_list): + # # Transfer input to GPU + # f_gpu = cp.asarray(f, dtype=cp.float32) + + # # Perform FFT convolution on GPU + # p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same') + + # # Scale by grid spacing (on GPU) + # p_gpu *= self._gpu_cache['gdx_prod'] + + # # Transfer result back to CPU + # results.append(cp.asnumpy(p_gpu)) + + # if verbose: + # print(f' - Solved {i+1}/{len(f_list)}') + + # if verbose: + # print('Poisson solver (GPU batch) finished') + # gtimer.end() + + # return results + + # elif self.stype == 'S': + # raise SyntaxError('Spatial solver not suggested for structured grids') + + # # ---------------------------------------------------------------------- + # # unstructured points - not optimized for GPU + # elif self.ptype == 'P': + # if verbose: + # print(' - GPU solver not optimized for unstructured points, using CPU') + # return [self.solve(f, verbose) for f in f_list] + + # # If we reach here, it's an unsupported case + # raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver") def clear_gpu_cache(self): """Clear GPU cache to free GPU memory"""