ykcheungab
52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
"""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
|