dq6a reach 18 ms per frame

This commit is contained in:
radxa-d6a
2026-04-03 12:52:48 +08:00
parent 424a4e37ce
commit 0a281c982c
4 changed files with 211 additions and 74 deletions
+40 -43
View File
@@ -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")