From 0a281c982cf93501dd00dfe86e04c0589ba2d23d Mon Sep 17 00:00:00 2001 From: radxa-d6a Date: Fri, 3 Apr 2026 12:52:48 +0800 Subject: [PATCH] dq6a reach 18 ms per frame --- benchmark.py | 17 ++++- src/orisys/HHD.py | 139 +++++++++++++++++++++++++++++++++++++---- src/orisys/Stitcher.py | 83 ++++++++++++------------ src/orisys/sensor.py | 46 ++++++++------ 4 files changed, 211 insertions(+), 74 deletions(-) diff --git a/benchmark.py b/benchmark.py index f448128..f37ec5a 100644 --- a/benchmark.py +++ b/benchmark.py @@ -12,15 +12,30 @@ if _local_src not in sys.path: import orisys import numpy as np import time +import argparse def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--video","-v", type=str, default="0", help="Camera ID / Video file path") + parser.add_argument("--config","-c", type=str, default="./config/ddjx01.json", help="Configuration file path") + parser.add_argument("--cal","-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path") + parser.add_argument("--verbose","-verbose", type=bool, default=True, help="Print detailed information") + parser.add_argument("--debug","-debug", type=bool, default=True, help="Print detailed information") + args = parser.parse_args() print(f"Using orisys module: {orisys.__file__}") # 步骤1：创建传感器对象 (必要步骤)， # config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径） # cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件 # 方式1: 从视频文件读取 - sensor = orisys.Sensor(2, config_name="./config/ddjx01_q6a.json", cal_path= "./config/ddjx01.npy", isstitch=True, cuda=False, verbose=True, backend='cpu') + try: + input = int(args.video) + is_camera = True + sensor = orisys.Sensor(input, config_name=args.config, cal_path= args.cal, verbose=args.verbose, debug=args.debug) + except: + is_camera = False + sensor = orisys.Sensor(args.video, config_name=args.config, cal_path= args.cal, verbose=args.verbose, debug=args.debug) + # sensor = orisys.Sensor(2, config_name="./config/ddjx01_q6a.json", cal_path= "./config/ddjx01.npy", isstitch=True, cuda=False, verbose=True, backend='cpu') # 方式2: 从摄像头读取（取消注释使用） # sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py index 3c6f919..0a2988e 100644 --- a/src/orisys/HHD.py +++ b/src/orisys/HHD.py @@ -26,12 +26,107 @@ import timeit import sys # import cv2 from scipy import signal, spatial +from numba import njit from .backend_availability import ( CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE, cp, cp_signal, _torque_kernel, cl, cla, cl_rfftn, cl_irfftn ) + + +# -- Numba-accelerated finite-difference kernels for 2D HHD ------------------ +# Matches numpy.gradient's second-order boundary stencil exactly. + +@njit(cache=True) +def _divcurl_2d_nb(vfield, dy, dx, div_out, curl_out): + """Fused divergence + curl of a 2D vector field via finite differences.""" + H = vfield.shape[0] + W = vfield.shape[1] + idy2 = 1.0 / (2.0 * dy) + idx2 = 1.0 / (2.0 * dx) + idy_edge = 1.0 / (2.0 * dy) + idx_edge = 1.0 / (2.0 * dx) + for y in range(H): + for x in range(W): + # du/dx + if x == 0: + dudx = (-3.0*vfield[y,0,0] + 4.0*vfield[y,1,0] - vfield[y,2,0]) * idx_edge + elif x == W - 1: + dudx = (vfield[y,W-3,0] - 4.0*vfield[y,W-2,0] + 3.0*vfield[y,W-1,0]) * idx_edge + else: + dudx = (vfield[y,x+1,0] - vfield[y,x-1,0]) * idx2 + # dv/dy + if y == 0: + dvdy = (-3.0*vfield[0,x,1] + 4.0*vfield[1,x,1] - vfield[2,x,1]) * idy_edge + elif y == H - 1: + dvdy = (vfield[H-3,x,1] - 4.0*vfield[H-2,x,1] + 3.0*vfield[H-1,x,1]) * idy_edge + else: + dvdy = (vfield[y+1,x,1] - vfield[y-1,x,1]) * idy2 + # dv/dx + if x == 0: + dvdx = (-3.0*vfield[y,0,1] + 4.0*vfield[y,1,1] - vfield[y,2,1]) * idx_edge + elif x == W - 1: + dvdx = (vfield[y,W-3,1] - 4.0*vfield[y,W-2,1] + 3.0*vfield[y,W-1,1]) * idx_edge + else: + dvdx = (vfield[y,x+1,1] - vfield[y,x-1,1]) * idx2 + # du/dy + if y == 0: + dudy = (-3.0*vfield[0,x,0] + 4.0*vfield[1,x,0] - vfield[2,x,0]) * idy_edge + elif y == H - 1: + dudy = (vfield[H-3,x,0] - 4.0*vfield[H-2,x,0] + 3.0*vfield[H-1,x,0]) * idy_edge + else: + dudy = (vfield[y+1,x,0] - vfield[y-1,x,0]) * idy2 + + div_out[y, x] = dudx + dvdy + curl_out[y, x] = dvdx - dudy + + +@njit(cache=True) +def _grad_rotgrad_h_2d_nb(nD, nRu, vfield, dy, dx, d_out, r_out, h_out): + """Fused gradient(nD) + rotated_gradient(nRu) + h = vfield - d - r.""" + H = nD.shape[0] + W = nD.shape[1] + idy2 = 1.0 / (2.0 * dy) + idx2 = 1.0 / (2.0 * dx) + idy_edge = 1.0 / (2.0 * dy) + idx_edge = 1.0 / (2.0 * dx) + for y in range(H): + for x in range(W): + # gradient(nD) → d = (ddx_nD, ddy_nD) + if x == 0: + ddx_nD = (-3.0*nD[y,0] + 4.0*nD[y,1] - nD[y,2]) * idx_edge + ddx_nRu = (-3.0*nRu[y,0] + 4.0*nRu[y,1] - nRu[y,2]) * idx_edge + elif x == W - 1: + ddx_nD = (nD[y,W-3] - 4.0*nD[y,W-2] + 3.0*nD[y,W-1]) * idx_edge + ddx_nRu = (nRu[y,W-3] - 4.0*nRu[y,W-2] + 3.0*nRu[y,W-1]) * idx_edge + else: + ddx_nD = (nD[y,x+1] - nD[y,x-1]) * idx2 + ddx_nRu = (nRu[y,x+1] - nRu[y,x-1]) * idx2 + + if y == 0: + ddy_nD = (-3.0*nD[0,x] + 4.0*nD[1,x] - nD[2,x]) * idy_edge + ddy_nRu = (-3.0*nRu[0,x] + 4.0*nRu[1,x] - nRu[2,x]) * idy_edge + elif y == H - 1: + ddy_nD = (nD[H-3,x] - 4.0*nD[H-2,x] + 3.0*nD[H-1,x]) * idy_edge + ddy_nRu = (nRu[H-3,x] - 4.0*nRu[H-2,x] + 3.0*nRu[H-1,x]) * idy_edge + else: + ddy_nD = (nD[y+1,x] - nD[y-1,x]) * idy2 + ddy_nRu = (nRu[y+1,x] - nRu[y-1,x]) * idy2 + + d0 = ddx_nD + d1 = ddy_nD + r0 = -ddy_nRu + r1 = ddx_nRu + + d_out[y, x, 0] = d0 + d_out[y, x, 1] = d1 + r_out[y, x, 0] = r0 + r_out[y, x, 1] = r1 + h_out[y, x, 0] = vfield[y, x, 0] - d0 - r0 + h_out[y, x, 1] = vfield[y, x, 1] - d1 - r1 + + # --------------------------------------------------------------------------- class nHHD(object): @@ -79,8 +174,23 @@ class nHHD(object): raise ValueError("nHHD Solver works only for 2 and 3 dimensions") self._force_cpu_cache = None + self._decompose_bufs = None self.psolver.prepare(False) + def _ensure_decompose_bufs(self, shape2d): + """Lazily allocate reusable buffers for the 2D decompose fast path.""" + if self._decompose_bufs is not None and self._decompose_bufs['shape'] == shape2d: + return self._decompose_bufs + self._decompose_bufs = { + 'shape': shape2d, + 'div': numpy.empty(shape2d, dtype=numpy.float32), + 'curl': numpy.empty(shape2d, dtype=numpy.float32), + 'd': numpy.empty(shape2d + (2,), dtype=numpy.float32), + 'r': numpy.empty(shape2d + (2,), dtype=numpy.float32), + 'h': numpy.empty(shape2d + (2,), dtype=numpy.float32), + } + return self._decompose_bufs + def get_force(self, vfield, verbose=False): if self.poisson_backend == 'cupy': return self.get_force_cuda(vfield, verbose) @@ -129,19 +239,19 @@ class nHHD(object): print( '\n ==> nhhd.decompose', vfield.shape) if self.dim == 2: + shape2d = (vfield.shape[0], vfield.shape[1]) + bufs = self._ensure_decompose_bufs(shape2d) + dy, dx = self.mesh.dx[0], self.mesh.dx[1] - # compute div curl - (self.div, self.curlw) = self.mesh.divcurl(vfield) - - if self.poisson_backend == 'opencl': - self.div = numpy.ascontiguousarray(self.div, dtype=numpy.float32) - self.curlw = numpy.ascontiguousarray(self.curlw, dtype=numpy.float32) + vf32 = vfield if vfield.dtype == numpy.float32 else vfield.astype(numpy.float32) + _divcurl_2d_nb(vf32, dy, dx, bufs['div'], bufs['curl']) + self.div = bufs['div'] + self.curlw = bufs['curl'] if verbose > 0: print( 'div =', self.div.shape, self.div.min(), self.div.max()) print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) - # compute potentials if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE): self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=self.debug) else: @@ -151,14 +261,19 @@ class nHHD(object): print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max()) - # compute fields as gradients of potentials - self.d = self.mesh.gradient(self.nD) + nD_f32 = self.nD if self.nD.dtype == numpy.float32 else self.nD.astype(numpy.float32) + nRu_f32 = self.nRu if self.nRu.dtype == numpy.float32 else self.nRu.astype(numpy.float32) + _grad_rotgrad_h_2d_nb(nD_f32, nRu_f32, vf32, dy, dx, + bufs['d'], bufs['r'], bufs['h']) + self.d = bufs['d'] + self.r = bufs['r'] + self.h = bufs['h'] + if verbose > 0: print( 'd =', self.d.shape, self.d.min(), self.d.max()) - - self.r = self.mesh.rotated_gradient(self.nRu) - if verbose > 0: print( 'r =', self.r.shape, self.r.min(), self.r.max()) + print( 'h =', self.h.shape, self.h.min(), self.h.max()) + return else: # compute div curl diff --git a/src/orisys/Stitcher.py b/src/orisys/Stitcher.py index e6e4cfd..f74e18a 100644 --- a/src/orisys/Stitcher.py +++ b/src/orisys/Stitcher.py @@ -3,6 +3,35 @@ import numpy as np import scipy.fft from scipy import ndimage import cv2 +from numba import njit + +@njit(cache=True) +def _blend_kernel(blend_patch, wsum_patch, cov_patch, src_patch, w_patch): + """Fused incremental weighted-average composite (numba-accelerated).""" + sh = src_patch.shape[0] + sw = src_patch.shape[1] + for y in range(sh): + for x in range(sw): + w = w_patch[y, x] + s0 = np.float32(src_patch[y, x, 0]) + s1 = np.float32(src_patch[y, x, 1]) + s2 = np.float32(src_patch[y, x, 2]) + if cov_patch[y, x] == 0: + blend_patch[y, x, 0] = s0 + blend_patch[y, x, 1] = s1 + blend_patch[y, x, 2] = s2 + wsum_patch[y, x] = w + else: + ws = wsum_patch[y, x] + w + if ws < 1e-6: + ws = 1e-6 + a = w / ws + blend_patch[y, x, 0] += (s0 - blend_patch[y, x, 0]) * a + blend_patch[y, x, 1] += (s1 - blend_patch[y, x, 1]) * a + blend_patch[y, x, 2] += (s2 - blend_patch[y, x, 2]) * a + wsum_patch[y, x] = ws + cov_patch[y, x] = 1 + class Stitcher: @@ -407,7 +436,7 @@ class Stitcher: return wy[:, None] * wx[None, :] def _blend_tile_inplace(self, blend, wsum, cov, src, weight, y0, x0, scratch=None): - """Copy non-overlap tiles directly; blend only when this tile touches existing content.""" + """Composite a tile onto the canvas using numba-fused weighted average.""" h, w = src.shape[:2] H, W = blend.shape[:2] y0d, x0d = max(0, y0), max(0, x0) @@ -417,36 +446,15 @@ class Stitcher: sy0, sx0 = y0d - y0, x0d - x0 sh, sw = y1d - y0d, x1d - x0d - w_patch = weight[sy0 : sy0 + sh, sx0 : sx0 + sw] - blend_patch = blend[y0d:y1d, x0d:x1d] - wsum_patch = wsum[y0d:y1d, x0d:x1d] - cov_patch = cov[y0d:y1d, x0d:x1d] - src_patch = src[sy0 : sy0 + sh, sx0 : sx0 + sw] - if not np.any(cov_patch): - blend_patch[:] = src_patch - wsum_patch[:] = w_patch - elif scratch is not None: - s_src, s_alpha = scratch - sf = s_src[:sh, :sw] - sa = s_alpha[:sh, :sw] - sf[:] = src_patch - np.add(wsum_patch, w_patch, out=wsum_patch) - np.maximum(wsum_patch, 1e-6, out=sa) - np.divide(w_patch, sa, out=sa) - np.subtract(sf, blend_patch, out=sf) - sf[:, :, 0] *= sa - sf[:, :, 1] *= sa - sf[:, :, 2] *= sa - np.add(blend_patch, sf, out=blend_patch) - else: - total = wsum_patch + w_patch - alpha = w_patch / np.maximum(total, 1e-6) - src_patch_f = src_patch.astype(np.float32, copy=False) - blend_patch += (src_patch_f - blend_patch) * alpha[:, :, None] - wsum_patch[:] = total - - cov_patch[:] = 1 + src_f = src[sy0:sy0+sh, sx0:sx0+sw].astype(np.float32, copy=False) + _blend_kernel( + blend[y0d:y1d, x0d:x1d], + wsum[y0d:y1d, x0d:x1d], + cov[y0d:y1d, x0d:x1d], + src_f, + weight[sy0:sy0+sh, sx0:sx0+sw], + ) def _mean_gray_bgr(self, bgr): ch = cv2.mean(bgr) @@ -641,17 +649,6 @@ class Stitcher: t_after_reset = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 blend_margin = 16 if blend else 0 - max_th = max(h0c, *(tb.shape[0] for tb in tiles_bgr_c)) if tiles_bgr_c else h0c - max_tw = max(w0c, *(tb.shape[1] for tb in tiles_bgr_c)) if tiles_bgr_c else w0c - _ss = canvas_cache.get("scratch_src") - _sa = canvas_cache.get("scratch_alpha") - if _ss is None or _ss.shape[0] < max_th or _ss.shape[1] < max_tw: - _ss = np.empty((max_th, max_tw, 3), dtype=np.float32) - canvas_cache["scratch_src"] = _ss - if _sa is None or _sa.shape[0] < max_th or _sa.shape[1] < max_tw: - _sa = np.empty((max_th, max_tw), dtype=np.float32) - canvas_cache["scratch_alpha"] = _sa - _scratch = (_ss, _sa) blended_tiles = 0 for tile_bgr, H_use, paste in zip(tiles_bgr_c, H_use_list, pastes): @@ -662,13 +659,13 @@ class Stitcher: wkey = (th, tw, blend_margin) if wkey not in weight_maps: weight_maps[wkey] = self._tile_weight_map(th, tw, blend_margin) - self._blend_tile_inplace(blend_buf, wsum, cov, tile_bgr, weight_maps[wkey], y0, x0, _scratch) + self._blend_tile_inplace(blend_buf, wsum, cov, tile_bgr, weight_maps[wkey], y0, x0) blended_tiles += 1 ref_wkey = (h0c, w0c, blend_margin) if ref_wkey not in weight_maps: weight_maps[ref_wkey] = self._tile_weight_map(h0c, w0c, blend_margin) - self._blend_tile_inplace(blend_buf, wsum, cov, t_ref_bgr, weight_maps[ref_wkey], ref_y0, ref_x0, _scratch) + self._blend_tile_inplace(blend_buf, wsum, cov, t_ref_bgr, weight_maps[ref_wkey], ref_y0, ref_x0) t_after_blend = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 crop_rect = canvas_cache.get("crop_rect") diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index ba54a9a..10a01c6 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -114,6 +114,8 @@ class Sensor: # self.stitch_calibration = [] self.frame = self.img = None self.initial_img = None + self._img_bgr_raw = None + self._img_gray_buf = None shape = (self.__config_cam["resolution_stitch"][1], self.__config_cam["resolution_stitch"][0]) self.img_size_raw = None self.img_filtered = None @@ -212,35 +214,33 @@ class Sensor: # Start processing timing t_processing_start = time.perf_counter() - # self.frame = frame.copy() self.frame = frame - # Extract and stitch tiles - # self.stitcher.get_stitch_unit(frame) self._extract_tiles(frame) t_after_extract = time.perf_counter() - img = np.ascontiguousarray(self.stitcher.stitch_rt_optimized(blend=True)) + img_bgr = self.stitcher.stitch_rt_optimized(blend=False) t_after_stitch = time.perf_counter() - self.img_size_raw = img.shape - # Resize - img = cv2.resize(img, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_LINEAR) - t_after_resize = time.perf_counter() - - if self.mirror: - img = cv2.flip(img, 1) - t_after_mirror = time.perf_counter() + self.img_size_raw = img_bgr.shape + self._img_bgr_raw = img_bgr - self.img = img.astype(np.uint8) - t_after_astype = time.perf_counter() + gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) + res = tuple(self._config_resolution_stitch) + if self._img_gray_buf is None or self._img_gray_buf.shape != (res[1], res[0]): + self._img_gray_buf = np.empty((res[1], res[0]), dtype=np.uint8) + cv2.resize(gray, res, dst=self._img_gray_buf, interpolation=cv2.INTER_LINEAR) + t_after_resize = time.perf_counter() + + if self.mirror: + cv2.flip(self._img_gray_buf, 1, dst=self._img_gray_buf) + + self.img = self._img_gray_buf t_processing_end = time.perf_counter() self._last_processing_time = t_processing_end - t_processing_start self._last_stage_times["extract_tiles"] = t_after_extract - t_processing_start self._last_stage_times["stitch"] = t_after_stitch - t_after_extract - self._last_stage_times["resize"] = t_after_resize - t_after_stitch - self._last_stage_times["mirror"] = t_after_mirror - t_after_resize - self._last_stage_times["astype"] = t_after_astype - t_after_mirror + self._last_stage_times["resize"] = t_processing_end - t_after_stitch # FPS calculation self._fps = 1.0 / (time.perf_counter() - self.last_log_time) @@ -248,6 +248,16 @@ class Sensor: return self.img.copy() + def _get_img_bgr(self): + """Produce a resized+mirrored BGR image on demand (for display).""" + if self._img_bgr_raw is None: + return None + res = tuple(self._config_resolution_stitch) + bgr = cv2.resize(self._img_bgr_raw, res, interpolation=cv2.INTER_LINEAR) + if self.mirror: + cv2.flip(bgr, 1, dst=bgr) + return bgr + def detect_motion(self, threshold=6): # detect motion by optical flow magnitude @@ -410,7 +420,7 @@ class Sensor: else None ), self.info.CONTOUR: lambda: self.contour.copy() if self.contour is not None else np.array([]), - self.info.IMG: lambda: self.img.copy() if self.img is not None else np.array([]), + self.info.IMG: lambda: self._get_img_bgr() if self._img_bgr_raw is not None else np.array([]), } return tuple(mapping.get(arg, lambda: None)() for arg in args if isinstance(arg, self.info))