fix numba stiching blend error and add numpy blend fallback
This commit is contained in:
@@ -0,0 +1,51 @@
|
||||
"""Fused stitch blend kernel (Numba JIT when available)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from numba import njit
|
||||
|
||||
_NUMBA_AVAILABLE = True
|
||||
except ImportError:
|
||||
_NUMBA_AVAILABLE = False
|
||||
|
||||
def njit(*args, **kwargs):
|
||||
def decorator(fn):
|
||||
return fn
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
if _NUMBA_AVAILABLE:
|
||||
|
||||
@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
|
||||
|
||||
else:
|
||||
blend_kernel = None
|
||||
Reference in New Issue
Block a user