From ab97efe0398c91f8df29c5c9fbce9f8ee776e737 Mon Sep 17 00:00:00 2001 From: ykcheungab Date: Sun, 23 Aug 2026 16:24:53 +0800 Subject: [PATCH] refractor stitcher class into module --- src/orisys/Stitcher.py | 1550 +------------------ src/orisys/stitch/__init__.py | 33 + src/orisys/stitch/api.py | 324 ++++ src/orisys/stitch/backends/__init__.py | 47 + src/orisys/stitch/backends/base.py | 29 + src/orisys/stitch/backends/numba_backend.py | 53 + src/orisys/stitch/backends/numpy_backend.py | 61 + src/orisys/stitch/boxes.py | 217 +++ src/orisys/stitch/config.py | 76 + src/orisys/stitch/feature.py | 599 +++++++ src/orisys/stitch/ops.py | 136 ++ src/orisys/stitch/tiles.py | 86 + tests/test_stitch.py | 168 ++ 13 files changed, 1845 insertions(+), 1534 deletions(-) create mode 100644 src/orisys/stitch/__init__.py create mode 100644 src/orisys/stitch/api.py create mode 100644 src/orisys/stitch/backends/__init__.py create mode 100644 src/orisys/stitch/backends/base.py create mode 100644 src/orisys/stitch/backends/numba_backend.py create mode 100644 src/orisys/stitch/backends/numpy_backend.py create mode 100644 src/orisys/stitch/boxes.py create mode 100644 src/orisys/stitch/config.py create mode 100644 src/orisys/stitch/feature.py create mode 100644 src/orisys/stitch/ops.py create mode 100644 src/orisys/stitch/tiles.py create mode 100644 tests/test_stitch.py diff --git a/src/orisys/Stitcher.py b/src/orisys/Stitcher.py index 02a2719..f3444d7 100644 --- a/src/orisys/Stitcher.py +++ b/src/orisys/Stitcher.py @@ -1,1539 +1,21 @@ -import os -import numpy as np -import scipy.fft -from scipy import ndimage -import cv2 +"""Compatibility shim: prefer ``orisys.stitch.Stitcher``. + +Historical import path ``from orisys.Stitcher import Stitcher`` remains supported. +""" + +from .stitch import Stitcher +from .stitch.backends.numpy_backend import blend_kernel_numpy as _blend_kernel_numpy +from .stitch.config import BLEND_MODES as _BLEND_MODES +from .stitch.config import DEFAULT_BLEND_MODE as _DEFAULT_BLEND_MODE +from .stitch.config import numba_blend_available, resolve_blend_mode as _resolve_blend_mode + +# Backward-compatible aliases used by older call sites / tests. +_blend_kernel = _blend_kernel_numpy +_NUMBA_AVAILABLE = numba_blend_available() try: - from orisys_blend.kernel import _NUMBA_AVAILABLE, blend_kernel as _blend_kernel_numba + from orisys_blend.kernel import blend_kernel as _blend_kernel_numba except Exception: - _NUMBA_AVAILABLE = False _blend_kernel_numba = None -# Re-export availability from the pure-Python kernel package when import succeeded -# but numba itself was missing (blend_kernel is None). -if _blend_kernel_numba is None: - _NUMBA_AVAILABLE = False - - -def _blend_kernel_numpy(blend_patch, wsum_patch, cov_patch, src_patch, w_patch): - """Incremental weighted-average composite (pure NumPy, no JIT). - - Uses contiguous arithmetic instead of boolean fancy indexing. Callers must - zero ``blend`` / ``wsum`` on uncovered pixels (as the stitcher reset does); - then the covered lerp also paints first-touch pixels correctly when ``w > 0``. - """ - w = w_patch if w_patch.dtype == np.float32 else np.asarray(w_patch, dtype=np.float32) - src = src_patch if src_patch.dtype == np.float32 else np.asarray(src_patch, dtype=np.float32) - - uncovered = cov_patch == 0 - if uncovered.all(): - # First paint over a cleared patch: direct copy. - np.copyto(blend_patch, src) - np.copyto(wsum_patch, w) - cov_patch.fill(1) - return - - # Contiguous lerp. With zeroed uncovered pixels this also equals a first paint - # when w > 0 (a ≈ 1 → blend becomes src). - ws = wsum_patch + w - a = w / np.maximum(ws, np.float32(1e-6)) - # Channel loop keeps temps HxW instead of one HxWx3 gather/scatter. - for c in range(src.shape[2]): - bc = blend_patch[..., c] - bc += (src[..., c] - bc) * a - np.copyto(wsum_patch, ws) - - # Match numba on rare first-touch pixels with ~0 weight (lerp would leave 0). - if uncovered.any(): - need_fix = uncovered & (w < np.float32(1e-6)) - if need_fix.any(): - blend_patch[need_fix] = src[need_fix] - wsum_patch[need_fix] = w[need_fix] - - cov_patch.fill(1) - - -# Backward-compatible alias used by older call sites / tests. -_blend_kernel = _blend_kernel_numpy - -# JSON ``stitcher.blend`` / ``blend_list``: none = fast paste; numpy/numba = feathered. -_BLEND_MODES = frozenset({"none", "numpy", "numba"}) -_DEFAULT_BLEND_MODE = "none" - - -def _resolve_blend_mode(config_stitcher=None): - """Resolve stitcher blend mode from env, then JSON ``blend``. - - Returns ``(mode, fell_back_from_numba)`` where mode is one of - ``none`` / ``numpy`` / ``numba``. - - Priority: ``ORISYS_STITCH_BLEND``, ``ORISYS_STITCH_BLEND_BACKEND``, - then ``config_stitcher["blend"]`` / legacy ``blend_backend``, then ``none``. - """ - requested = None - for key in ("ORISYS_STITCH_BLEND", "ORISYS_STITCH_BLEND_BACKEND"): - env = os.environ.get(key, "").strip().lower() - if env in _BLEND_MODES: - requested = env - break - - if requested is None and isinstance(config_stitcher, dict): - for key in ("blend", "blend_backend"): - cfg = str(config_stitcher.get(key, "")).strip().lower() - if cfg in _BLEND_MODES: - requested = cfg - break - - if requested is None: - requested = _DEFAULT_BLEND_MODE - - if requested == "numba" and not _NUMBA_AVAILABLE: - return "numpy", True - return requested, False - - -class Stitcher: - - def __init__(self,isstitch=False, config_stitcher=None,verbose=False): - - # Settings - self.isstitch = isstitch - self.config_stitcher = config_stitcher - self.verbose = verbose - # Find Boxes Related - self.min_area_frac = config_stitcher["min_area_frac"] - self.max_area_frac = config_stitcher["max_area_frac"] - self.max_aspect_ratio = config_stitcher["max_aspect_ratio"] - - # stitch related - self.stitch_valid = False - self.boxes = self.units = None - self.index = config_stitcher["index"] - self.stitch_calibration = self.tile_list = [] - self.ncol = config_stitcher["n"][0] - self.nrow = config_stitcher["n"][1] - self.nominal_ov_pos=config_stitcher["nominal_ov_pos"] - self.margin=config_stitcher["margin"] - self.bound=self.nominal_ov_pos - self.margin - self.row_tol = config_stitcher["row_tol"] - self.box_size = config_stitcher["box_size"] - - # stich related (BF-matching) - self._blend_coef_cache = {} - self._feature_stitch_cache = { - "ref_idx": None, - "moving_indices": None, - "tile_shapes": None, - "transforms": None, - "canvas": {}, - } - self._feature_stitch_frame_idx = 0 - self._feature_stitch_print_every = 30 - self._feature_stitch_pad = 5 - self._feature_stitch_match_brightness = False - self._blend_mode, fell_back = _resolve_blend_mode(config_stitcher) - # Feather path uses numpy/numba; ``none`` keeps a numpy fallback if forced on. - self._blend_backend = "numpy" if self._blend_mode == "none" else self._blend_mode - self._numba_failed = False - - if fell_back and self.isstitch: - print( - "Warning: stitcher blend=numba requested but numba is " - "unavailable; using numpy blend path." - ) - elif self.isstitch and self.verbose: - print(f"Stitch blend: {self._blend_mode}") - - if self.verbose: - print(f"Stitcher initialized with config: {self.config_stitcher}") - - @property - def blend_mode(self): - """Configured blend mode: ``none``, ``numpy``, or ``numba``.""" - return self._blend_mode - - def configure_feature_stitch( - self, - *, - match_brightness=None, - pad=None, - print_every=None, - blend=None, - blend_backend=None, - ): - """Update runtime options for feature-based stitching. - - ``blend`` accepts ``none`` / ``numpy`` / ``numba`` (same as JSON ``stitcher.blend``). - ``blend_backend`` is a legacy alias for the same values. - """ - if match_brightness is not None: - self._feature_stitch_match_brightness = bool(match_brightness) - if pad is not None: - self._feature_stitch_pad = max(0, int(pad)) - if print_every is not None: - self._feature_stitch_print_every = max(1, int(print_every)) - mode = blend if blend is not None else blend_backend - if mode is not None: - mode = str(mode).strip().lower() - if mode not in _BLEND_MODES: - raise ValueError( - f"Unknown blend={mode!r}; expected one of {sorted(_BLEND_MODES)}" - ) - if mode == "numba" and (not _NUMBA_AVAILABLE or self._numba_failed): - print("Warning: numba unavailable; using numpy blend.") - mode = "numpy" - self._blend_mode = mode - self._blend_backend = "numpy" if mode == "none" else mode - if mode == "numba": - self._numba_failed = False - - def _resolve_call_blend(self, blend): - """Map a ``stitch_rt_optimized(blend=...)`` arg to ``(use_feather, backend)``.""" - if blend is None: - mode = self._blend_mode - elif isinstance(blend, bool): - if blend: - mode = self._blend_mode if self._blend_mode != "none" else "numpy" - else: - return False, self._blend_backend - else: - mode = str(blend).strip().lower() - if mode not in _BLEND_MODES: - raise ValueError( - f"Unknown blend={mode!r}; expected one of {sorted(_BLEND_MODES)}" - ) - if mode == "numba" and not _NUMBA_AVAILABLE: - mode = "numpy" - - if mode == "numba" and self._numba_failed: - mode = "numpy" - - if mode == "none": - return False, self._blend_backend - return True, mode - - def _mark_numba_failed(self, exc): - """Stick to numpy after a Numba JIT/runtime failure (log once).""" - if self._numba_failed: - return - self._numba_failed = True - self._blend_backend = "numpy" - if self._blend_mode == "numba": - self._blend_mode = "numpy" - print( - f"Warning: numba blend failed ({type(exc).__name__}: {exc}); " - f"falling back to numpy." - ) - - def _get_blend_coef(self, ov_pos, dim, ndim): - key = (ov_pos, dim, ndim) - coef = self._blend_coef_cache.get(key) - if coef is not None: - return coef - blend_coef = np.linspace(1.0, 0.0, ov_pos, dtype=np.float32) - if dim == 1: - blend_coef = blend_coef.reshape(1, ov_pos, 1) if ndim == 3 else blend_coef.reshape(1, ov_pos) - else: - blend_coef = blend_coef.reshape(ov_pos, 1, 1) if ndim == 3 else blend_coef.reshape(ov_pos, 1) - self._blend_coef_cache[key] = blend_coef - return blend_coef - - def stitch_rt_optimized(self, change_sv=False, s_mean=28., v_mean=180., blend=None): - """Stitch the current tile list. - - ``blend`` overrides JSON/env for this call: - - ``None`` (default): use configured ``stitcher.blend`` - - ``True`` / ``False``: force feathered blend on/off - - ``\"none\"`` / ``\"numpy\"`` / ``\"numba\"``: force that mode for this call - """ - use_feather, blend_backend = self._resolve_call_blend(blend) - prev_backend = self._blend_backend - self._blend_backend = blend_backend - try: - feature_stitched = self.stitch_feature_optimized( - blend=use_feather, - match_brightness=self._feature_stitch_match_brightness, - pad=self._feature_stitch_pad, - print_every=self._feature_stitch_print_every, - ) - finally: - # Keep sticky numpy fallback after a Numba failure; do not restore numba. - if self._numba_failed: - self._blend_backend = "numpy" - else: - self._blend_backend = prev_backend - if feature_stitched is not None: - return feature_stitched - - # tile_list = self.tile_list - # stitch_calibration = self.stitch_calibration - # ncol = self.ncol - # nrow = self.nrow - - # img_stitched_height = None - - # for i in range(nrow): - # # Start with first tile of the row - # img_stitched_width = tile_list[i * ncol] - - # # Stitch horizontally within the row - # for j in range(1, ncol): - # img_1 = img_stitched_width - # img_2 = tile_list[i * ncol + j] - - # ov_pos = int(stitch_calibration[i * ncol + j - 1][1]) - # ov_pos, ov1, ov2 = self.__find_ov(img_1, img_2, ov_pos, dim=1) - - # img_stitched_width = self.__stitch_optimized( - # img_1, img_2, ov_pos, ov1, ov2, dim=1, blend=blend - # ) - - # # Stitch rows vertically - # if i == 0: - # img_stitched_height = img_stitched_width - # else: - # width_h = img_stitched_height.shape[1] - # width_w = img_stitched_width.shape[1] - # wdiff = width_h - width_w - - # if wdiff != 0: - # if wdiff > 0: - # padding = ( - # ((0, 0), (0, wdiff), (0, 0)) - # if len(img_stitched_width.shape) == 3 - # else ((0, 0), (0, wdiff)) - # ) - # img_stitched_width = np.pad(img_stitched_width, padding, mode='edge') - # else: - # img_stitched_width = img_stitched_width[:, :width_h] - - # ov_pos_h = int(stitch_calibration[(i + 1) * nrow - 1][1]) - # ov_pos_h, ov_h1, ov_h2 = self.__find_ov( - # img_stitched_height, img_stitched_width, ov_pos_h, dim=0 - # ) - - # img_stitched_height = self.__stitch_optimized( - # img_stitched_height, img_stitched_width, ov_pos_h, ov_h1, ov_h2, dim=0, blend=blend - # ) - - # return img_stitched_height - - def stitch_feature_optimized(self, blend=True, match_brightness=False, pad=5, print_every=30): - """ - Stitch the current tile grid by estimating tile-to-reference transforms once and - reusing them across frames. Falls back to overlap stitching when feature matching - cannot build a valid layout. - """ - plan = self._get_feature_stitch_plan() - if plan is None: - return None - - ref_idx, t_ref, moving_indices, moving_tiles, selected_positions = plan - self._feature_stitch_frame_idx += 1 - frame_idx = self._feature_stitch_frame_idx - tile_shapes = tuple(self.tile_list[pos].shape[:2] for pos in selected_positions) - cache = self._feature_stitch_cache - cache_changed = ( - cache["ref_idx"] != ref_idx - or cache["moving_indices"] != moving_indices - or cache["tile_shapes"] != tile_shapes - ) - - if cache_changed: - cache["ref_idx"] = ref_idx - cache["moving_indices"] = moving_indices - cache["tile_shapes"] = tile_shapes - cache["transforms"] = None - cache["canvas"] = {} - if self.verbose: - print( - "Feature stitch tile selection: " - f"ref_idx={ref_idx}, moving_indices={list(moving_indices)}" - ) - - H_use_list = cache["transforms"] - if H_use_list is None: - do_verbose_first = self.verbose and frame_idx % print_every == 0 - H_use_list = [] - for j, tile in enumerate(moving_tiles): - H_use_j, _ = self._estimate_bfmatcher_stitch_transform( - t_ref, - tile, - verbose=(do_verbose_first and j == 0), - ) - H_use_list.append(H_use_j) - cache["transforms"] = H_use_list - - if not any(H is not None for H in H_use_list): - return None - - return self._render_feature_tiles( - t_ref, - moving_tiles, - H_use_list, - pad=pad, - frame_idx=frame_idx, - print_every=print_every, - verbose=self.verbose, - match_brightness=match_brightness, - canvas_cache=cache["canvas"], - blend=blend, - ) - - # def __find_ov(self, img1, img2, ov_pos, dim=1): - - # _dim = len(img1.shape) - - # if dim == 1: - # if _dim == 3: - # # RGB case - # ov_1 = img1[:, -ov_pos:, :] - # ov_2 = img2[:, :ov_pos, :] - # else: - # # gray scale - # ov_1 = img1[:, -ov_pos:] - # ov_2 = img2[:, :ov_pos] - # elif dim == 0: - # if _dim == 3: - # ov_1 = img1[-ov_pos:, :, :] - # ov_2 = img2[:ov_pos, :, :] - # else: - # # gray scale - # ov_1 = img1[-ov_pos:, :] - # ov_2 = img2[:ov_pos, :] - - # return ov_pos, ov_1, ov_2 - - def _get_feature_stitch_plan(self): - if not self.isstitch: - return None - - if isinstance(self.index, int): - configured_indices = (self.index,) - elif isinstance(self.index, np.ndarray): - configured_indices = tuple(int(i) for i in self.index.tolist()) - else: - configured_indices = tuple(int(i) for i in self.index) - - total_tiles = len(configured_indices) - expected_tiles = self.nrow * self.ncol - if total_tiles <= 1 or total_tiles != expected_tiles or len(self.tile_list) < total_tiles: - return None - - # Use the configured grid order from self.index and keep the same reference tile - # location as the original stitch setup. - ref_pos = (self.nrow // 2) * self.ncol + (self.ncol // 2) - if ref_pos >= total_tiles: - return None - - selected_positions = tuple(range(total_tiles)) - ref_idx = configured_indices[ref_pos] - moving_positions = tuple(pos for pos in selected_positions if pos != ref_pos) - moving_indices = tuple(configured_indices[pos] for pos in moving_positions) - moving_tiles = [self.tile_list[pos] for pos in moving_positions] - return ref_idx, self.tile_list[ref_pos], moving_indices, moving_tiles, selected_positions - - def _tile_to_bgr(self, img): - if img.ndim == 2: - return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - return img - - def _tile_to_gray(self, bgr): - if bgr.ndim == 2: - return bgr - return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) - - def _ensure_u8_gray(self, gray): - if gray.dtype == np.uint8: - return gray - return cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) - - def _preprocess_gray_for_orb(self, g0, g1): - g0 = cv2.GaussianBlur(g0, (5, 5), 0) - g1 = cv2.GaussianBlur(g1, (5, 5), 0) - g0 = cv2.equalizeHist(g0) - g1 = cv2.equalizeHist(g1) - return g0, g1 - - def _estimate_bfmatcher_stitch_transform( - self, - t0, - t1, - *, - ransac_reproj_thresh=5.0, - ratio=0.9, - orb_nfeatures=5000, - verbose=True, - ): - """ - Estimate a 3x3 transform mapping tile coordinates into the reference tile frame. - """ - t0_bgr = self._tile_to_bgr(t0) - t1_bgr = self._tile_to_bgr(t1) - g0 = self._ensure_u8_gray(self._tile_to_gray(t0_bgr)) - g1 = self._ensure_u8_gray(self._tile_to_gray(t1_bgr)) - - if verbose: - print( - f"BF stitcher debug: t0 dtype={t0.dtype}, t1 dtype={t1.dtype}, " - f"g0 dtype={g0.dtype}, g0 range=({np.min(g0)}, {np.max(g0)}), " - f"g1 dtype={g1.dtype}, g1 range=({np.min(g1)}, {np.max(g1)})" - ) - - g0, g1 = self._preprocess_gray_for_orb(g0, g1) - orb = cv2.ORB_create( - nfeatures=orb_nfeatures, - scaleFactor=1.2, - nlevels=8, - edgeThreshold=15, - fastThreshold=5, - ) - k0, d0 = orb.detectAndCompute(g0, None) - k1, d1 = orb.detectAndCompute(g1, None) - - info = { - "ok": False, - "method": None, - "inliers": 0, - "mean_err": None, - "good_matches": 0, - "k0": k0, - "k1": k1, - "t0_bgr": t0_bgr, - "t1_bgr": t1_bgr, - } - if d0 is None or d1 is None or len(k0) < 4 or len(k1) < 4: - if verbose: - print( - "OpenCV BFMatcher stitcher: not enough features.", - f"k0={0 if k0 is None else len(k0)}", - f"k1={0 if k1 is None else len(k1)}", - f"d0_is_none={d0 is None}", - f"d1_is_none={d1 is None}", - ) - return None, info - - matches = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False).knnMatch(d0, d1, k=2) - good = [] - for m_n in matches: - if len(m_n) != 2: - continue - m, n = m_n - if m.distance < ratio * n.distance: - good.append(m) - - info["good_matches"] = len(good) - if verbose: - print(f"BFMatcher good matches: {len(good)}") - if len(good) < 4: - return None, info - - pts0 = np.float32([k0[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) - pts1 = np.float32([k1[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) - method = "affine_partial2d" - M, inlier_mask = cv2.estimateAffinePartial2D( - pts1, - pts0, - method=cv2.RANSAC, - ransacReprojThreshold=ransac_reproj_thresh, - ) - if M is not None: - H = np.vstack([M, [0, 0, 1]]).astype(np.float64) - else: - method = "homography" - H, inlier_mask = cv2.findHomography(pts1, pts0, cv2.RANSAC, ransac_reproj_thresh) - - info["method"] = method - info["inliers"] = int(inlier_mask.sum()) if inlier_mask is not None else 0 - if verbose: - print(f"Estimated {method} inliers: {info['inliers']}") - if H is None: - return None, info - - if abs(H[2, 2]) > 1e-12: - H = H / H[2, 2] - - mean_err = 1e9 - if inlier_mask is not None: - inlier_bool = inlier_mask.ravel().astype(bool) - if np.any(inlier_bool): - projected = cv2.perspectiveTransform(pts1[inlier_bool], H) - err = np.linalg.norm( - projected.reshape(-1, 2) - pts0[inlier_bool].reshape(-1, 2), - axis=1, - ) - mean_err = float(err.mean()) if err.size else 1e9 - info["mean_err"] = mean_err - info["ok"] = True - return H, info - - def _translation_from_H_to_ref(self, H): - if H is None: - return 0.0, 0.0 - H = np.asarray(H, dtype=np.float64) - if abs(H[2, 2]) > 1e-12: - H = H / H[2, 2] - pt = H @ np.array([[0.0], [0.0], [1.0]], dtype=np.float64) - return float(pt[0, 0]), float(pt[1, 0]) - - def _tile_weight_map(self, h, w, margin=16): - """ - Linear ramp weight that is 0 at tile edges and 1.0 at margin pixels inward. - Used for feathered blending at tile overlaps. - """ - wy = np.ones(h, dtype=np.float32) - wx = np.ones(w, dtype=np.float32) - m = min(margin, h // 2, w // 2) - if m > 0: - ramp = np.linspace(0.0, 1.0, m + 1, dtype=np.float32)[1:] - wy[:m] = ramp - wy[-m:] = ramp[::-1] - wx[:m] = ramp - wx[-m:] = ramp[::-1] - return wy[:, None] * wx[None, :] - - def _blend_tile_inplace(self, blend, wsum, cov, src, weight, y0, x0, scratch=None): - """Composite a tile onto the canvas using weighted average blending.""" - h, w = src.shape[:2] - H, W = blend.shape[:2] - y0d, x0d = max(0, y0), max(0, x0) - y1d, x1d = min(H, y0 + h), min(W, x0 + w) - if y0d >= y1d or x0d >= x1d: - return - - sy0, sx0 = y0d - y0, x0d - x0 - sh, sw = y1d - y0d, x1d - x0d - - src_f = src[sy0:sy0+sh, sx0:sx0+sw].astype(np.float32, copy=False) - args = ( - blend[y0d:y1d, x0d:x1d], - wsum[y0d:y1d, x0d:x1d], - cov[y0d:y1d, x0d:x1d], - src_f, - weight[sy0:sy0+sh, sx0:sx0+sw], - ) - if self._blend_backend == "numba" and _blend_kernel_numba is not None: - try: - _blend_kernel_numba(*args) - return - except Exception as exc: - self._mark_numba_failed(exc) - _blend_kernel_numpy(*args) - - def _mean_gray_bgr(self, bgr): - ch = cv2.mean(bgr) - return (ch[0] + ch[1] + ch[2]) / 3.0 - - def _match_bgr_brightness(self, tile_bgr, ref_mean_gray): - delta = int(round(ref_mean_gray - self._mean_gray_bgr(tile_bgr))) - if delta == 0: - return tile_bgr - lut = np.clip(np.arange(256, dtype=np.int16) + delta, 0, 255).astype(np.uint8) - return cv2.LUT(tile_bgr, lut) - - def _largest_fully_covered_rect(self, cov): - if cov.size == 0: - return None - - mask = cov > 0 - if not np.any(mask): - return None - - h, w = mask.shape - heights = np.zeros(w, dtype=np.int32) - best_area = 0 - best_rect = None - for y in range(h): - heights = np.where(mask[y], heights + 1, 0) - stack = [] - for x in range(w + 1): - cur_h = heights[x] if x < w else 0 - start = x - while stack and stack[-1][1] > cur_h: - start0, h0 = stack.pop() - area = h0 * (x - start0) - if area > best_area: - best_area = area - best_rect = (y - h0 + 1, y + 1, start0, x) - start = start0 - if not stack or stack[-1][1] < cur_h: - stack.append((start, cur_h)) - return best_rect - - def _compute_final_crop_rect(self, cov): - valid = cov > 0 - if not np.any(valid): - return None - - ys, xs = np.where(valid) - by0, by1 = int(ys.min()), int(ys.max()) + 1 - bx0, bx1 = int(xs.min()), int(xs.max()) + 1 - rect2 = self._largest_fully_covered_rect(cov[by0:by1, bx0:bx1]) - if rect2 is None: - return by0, by1, bx0, bx1 - - y0, y1, x0, x1 = rect2 - return by0 + y0, by0 + y1, bx0 + x0, bx0 + x1 - - def _geom_cache_key(self, t_ref_bgr, tiles_bgr_c, pad): - ref_h, ref_w = t_ref_bgr.shape[:2] - shapes = tuple(tb.shape[:2] for tb in tiles_bgr_c) - return ref_h, ref_w, shapes, pad - - def _ensure_feature_geom_cache( - self, - canvas_cache, - t_ref_bgr, - tiles_bgr_c, - H_use_list, - pad, - *, - allocate_blend_bufs=True, - ): - """Populate/reuse fixed-window geometry in canvas_cache. Returns geom_cache_hit.""" - h0c, w0c = t_ref_bgr.shape[:2] - geom_key = self._geom_cache_key(t_ref_bgr, tiles_bgr_c, pad) - geom_cache_hit = canvas_cache.get("geom_key") == geom_key - if geom_cache_hit: - if allocate_blend_bufs and ( - "blend_buf" not in canvas_cache - or "wsum_buf" not in canvas_cache - or "cov_buf" not in canvas_cache - ): - out_h = canvas_cache["out_h"] - out_w = canvas_cache["out_w"] - canvas_cache["blend_buf"] = np.zeros((out_h, out_w, 3), dtype=np.float32) - canvas_cache["wsum_buf"] = np.zeros((out_h, out_w), dtype=np.float32) - canvas_cache["cov_buf"] = np.zeros((out_h, out_w), dtype=np.uint8) - return True - - ref_corners = np.array([[0, 0], [w0c, 0], [w0c, h0c], [0, h0c]], dtype=np.float64) - pts = [ref_corners] - for tile_bgr, H_use in zip(tiles_bgr_c, H_use_list): - if H_use is None: - continue - hi, wi = tile_bgr.shape[:2] - tx, ty = self._translation_from_H_to_ref(H_use) - pts.append( - np.array( - [[tx, ty], [tx + wi, ty], [tx + wi, ty + hi], [tx, ty + hi]], - dtype=np.float64, - ) - ) - - all_xy = np.concatenate(pts, axis=0) - x_min_f = float(all_xy[:, 0].min()) - pad - y_min_f = float(all_xy[:, 1].min()) - pad - x_max_f = float(all_xy[:, 0].max()) + pad - y_max_f = float(all_xy[:, 1].max()) + pad - out_w = max(1, int(np.ceil(x_max_f - x_min_f))) - out_h = max(1, int(np.ceil(y_max_f - y_min_f))) - ref_y0 = int(np.floor(-y_min_f)) - ref_x0 = int(np.floor(-x_min_f)) - - pastes = [] - for H_use in H_use_list: - if H_use is None: - pastes.append(None) - continue - tx, ty = self._translation_from_H_to_ref(H_use) - pastes.append((int(np.floor(ty - y_min_f)), int(np.floor(tx - x_min_f)))) - - _dr_y0 = max(0, ref_y0) - _dr_x0 = max(0, ref_x0) - _dr_y1 = min(out_h, ref_y0 + h0c) - _dr_x1 = min(out_w, ref_x0 + w0c) - for _tb, _pa in zip(tiles_bgr_c, pastes): - if _pa is None: - continue - _py, _px = _pa - _th, _tw = _tb.shape[:2] - _dr_y0 = min(_dr_y0, max(0, _py)) - _dr_x0 = min(_dr_x0, max(0, _px)) - _dr_y1 = max(_dr_y1, min(out_h, _py + _th)) - _dr_x1 = max(_dr_x1, min(out_w, _px + _tw)) - - canvas_cache["geom_key"] = geom_key - canvas_cache["out_w"] = out_w - canvas_cache["out_h"] = out_h - canvas_cache["ref_y0"] = ref_y0 - canvas_cache["ref_x0"] = ref_x0 - canvas_cache["pastes"] = pastes - canvas_cache["x_min_f"] = x_min_f - canvas_cache["y_min_f"] = y_min_f - canvas_cache["crop_rect"] = None - canvas_cache["dirty_rect"] = (_dr_y0, _dr_y1, _dr_x0, _dr_x1) - canvas_cache["weight_maps"] = {} - canvas_cache["uint8_buf"] = np.zeros((out_h, out_w, 3), dtype=np.uint8) - if allocate_blend_bufs: - canvas_cache["blend_buf"] = np.zeros((out_h, out_w, 3), dtype=np.float32) - canvas_cache["wsum_buf"] = np.zeros((out_h, out_w), dtype=np.float32) - canvas_cache["cov_buf"] = np.zeros((out_h, out_w), dtype=np.uint8) - else: - # Coverage buffer only needed once to compute crop_rect for paste path. - canvas_cache["cov_buf"] = np.zeros((out_h, out_w), dtype=np.uint8) - canvas_cache.pop("blend_buf", None) - canvas_cache.pop("wsum_buf", None) - return False - - @staticmethod - def _paste_tile_clipped(canvas, tile, y0, x0, cov=None): - """Paste tile into canvas with clipping; optionally mark coverage.""" - th, tw = tile.shape[:2] - H, W = canvas.shape[:2] - y0d, x0d = max(0, y0), max(0, x0) - y1d, x1d = min(H, y0 + th), min(W, x0 + tw) - if y0d >= y1d or x0d >= x1d: - return - sy0, sx0 = y0d - y0, x0d - x0 - sh, sw = y1d - y0d, x1d - x0d - np.copyto(canvas[y0d:y1d, x0d:x1d], tile[sy0:sy0 + sh, sx0:sx0 + sw]) - if cov is not None: - cov[y0d:y1d, x0d:x1d] = 1 - - def _paste_feature_tiles_fast( - self, - t_ref_bgr, - tiles_bgr_c, - H_use_list, - canvas_cache, - *, - geom_cache_hit, - frame_idx=0, - print_every=30, - verbose=False, - t_profile0=0.0, - t_after_brightness=0.0, - t_after_geom=0.0, - ): - """Fixed-geometry uint8 paste path used when blend=False.""" - do_profile = verbose and frame_idx % print_every == 0 - out_w = canvas_cache["out_w"] - out_h = canvas_cache["out_h"] - ref_y0 = canvas_cache["ref_y0"] - ref_x0 = canvas_cache["ref_x0"] - pastes = canvas_cache["pastes"] - - canvas = canvas_cache.get("uint8_buf") - if canvas is None or canvas.shape[0] != out_h or canvas.shape[1] != out_w: - canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8) - canvas_cache["uint8_buf"] = canvas - - dirty = canvas_cache.get("dirty_rect") - if geom_cache_hit and dirty is not None: - dy0, dy1, dx0, dx1 = dirty - canvas[dy0:dy1, dx0:dx1] = 0 - elif not geom_cache_hit: - canvas.fill(0) - else: - canvas.fill(0) - t_after_reset = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - need_crop = canvas_cache.get("crop_rect") is None - cov = canvas_cache.get("cov_buf") if need_crop else None - if need_crop: - if cov is None or cov.shape != (out_h, out_w): - cov = np.zeros((out_h, out_w), dtype=np.uint8) - canvas_cache["cov_buf"] = cov - else: - if dirty is not None: - dy0, dy1, dx0, dx1 = dirty - cov[dy0:dy1, dx0:dx1] = 0 - else: - cov.fill(0) - - pasted_tiles = 0 - for tile_bgr, H_use, paste in zip(tiles_bgr_c, H_use_list, pastes): - if H_use is None or paste is None: - continue - y0, x0 = paste - self._paste_tile_clipped(canvas, tile_bgr, y0, x0, cov=cov) - pasted_tiles += 1 - self._paste_tile_clipped(canvas, t_ref_bgr, ref_y0, ref_x0, cov=cov) - t_after_paste = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - crop_rect = canvas_cache.get("crop_rect") - if crop_rect is None and cov is not None: - crop_rect = self._compute_final_crop_rect(cov) - canvas_cache["crop_rect"] = crop_rect - if crop_rect is not None: - y0, y1, x0, x1 = crop_rect - roi = canvas[y0:y1, x0:x1] - else: - roi = canvas - - rh, rw = roi.shape[:2] - stitch_buf = canvas_cache.get("stitch_buf") - if stitch_buf is None or stitch_buf.shape[0] != rh or stitch_buf.shape[1] != rw: - stitch_buf = np.empty((rh, rw, 3), dtype=np.uint8) - canvas_cache["stitch_buf"] = stitch_buf - np.copyto(stitch_buf, roi) - t_after_finalize = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - if do_profile: - xmf = canvas_cache.get("x_min_f", 0.0) - ymf = canvas_cache.get("y_min_f", 0.0) - print( - f"stitch_feature_optimized: canvas out_w={out_w}, out_h={out_h}, " - f"x_min={xmf:.2f}, y_min={ymf:.2f} (fast_paste)" - ) - print( - "stitch_feature_optimized steps: " - f"brightness={t_after_brightness - t_profile0:.6f}s, " - f"geom={t_after_geom - t_after_brightness:.6f}s ({'hit' if geom_cache_hit else 'miss'}), " - f"reset={t_after_reset - t_after_geom:.6f}s, " - f"paste={t_after_paste - t_after_reset:.6f}s ({pasted_tiles + 1} tiles incl ref), " - f"finalize={t_after_finalize - t_after_paste:.6f}s, " - f"total={t_after_finalize - t_profile0:.6f}s" - ) - return stitch_buf - - def _render_feature_tiles( - self, - t_ref_bgr, - tiles_bgr, - H_use_list, - *, - pad=5, - frame_idx=0, - print_every=30, - verbose=False, - match_brightness=False, - canvas_cache=None, - blend=True, - ): - """ - Composite multiple tiles into one canvas using translation only, where each tile - origin maps into the reference tile coordinates. - """ - if H_use_list is None: - return None - - if canvas_cache is None: - canvas_cache = {} - do_profile = verbose and frame_idx % print_every == 0 - t_profile0 = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - if t_ref_bgr.ndim == 2: - t_ref_bgr = cv2.cvtColor(t_ref_bgr, cv2.COLOR_GRAY2BGR) - h0c, w0c = t_ref_bgr.shape[:2] - if match_brightness: - ref_mean_gray = self._mean_gray_bgr(t_ref_bgr) - tiles_bgr_c = [ - self._match_bgr_brightness( - tb if tb.ndim == 3 else cv2.cvtColor(tb, cv2.COLOR_GRAY2BGR), - ref_mean_gray, - ) - for tb in tiles_bgr - ] - elif any(tb.ndim == 2 for tb in tiles_bgr): - tiles_bgr_c = [tb if tb.ndim == 3 else cv2.cvtColor(tb, cv2.COLOR_GRAY2BGR) for tb in tiles_bgr] - else: - tiles_bgr_c = tiles_bgr - t_after_brightness = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - geom_cache_hit = self._ensure_feature_geom_cache( - canvas_cache, - t_ref_bgr, - tiles_bgr_c, - H_use_list, - pad, - allocate_blend_bufs=bool(blend), - ) - t_after_geom = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - if not blend: - return self._paste_feature_tiles_fast( - t_ref_bgr, - tiles_bgr_c, - H_use_list, - canvas_cache, - geom_cache_hit=geom_cache_hit, - frame_idx=frame_idx, - print_every=print_every, - verbose=verbose, - t_profile0=t_profile0, - t_after_brightness=t_after_brightness, - t_after_geom=t_after_geom, - ) - - out_w = canvas_cache["out_w"] - out_h = canvas_cache["out_h"] - ref_y0 = canvas_cache["ref_y0"] - ref_x0 = canvas_cache["ref_x0"] - pastes = canvas_cache["pastes"] - - blend_buf = canvas_cache["blend_buf"] - wsum = canvas_cache["wsum_buf"] - cov = canvas_cache["cov_buf"] - weight_maps = canvas_cache["weight_maps"] - if geom_cache_hit: - _dr = canvas_cache.get("dirty_rect") - if _dr is not None: - _dy0, _dy1, _dx0, _dx1 = _dr - blend_buf[_dy0:_dy1, _dx0:_dx1] = 0 - wsum[_dy0:_dy1, _dx0:_dx1] = 0 - cov[_dy0:_dy1, _dx0:_dx1] = 0 - else: - blend_buf.fill(0) - wsum.fill(0) - cov.fill(0) - t_after_reset = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - blend_margin = 16 - - blended_tiles = 0 - for tile_bgr, H_use, paste in zip(tiles_bgr_c, H_use_list, pastes): - if H_use is None or paste is None: - continue - y0, x0 = paste - th, tw = tile_bgr.shape[:2] - 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) - 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) - t_after_blend = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - crop_rect = canvas_cache.get("crop_rect") - if crop_rect is None: - crop_rect = self._compute_final_crop_rect(cov) - canvas_cache["crop_rect"] = crop_rect - if crop_rect is not None: - y0, y1, x0, x1 = crop_rect - roi = blend_buf[y0:y1, x0:x1] - else: - roi = blend_buf - rh, rw = roi.shape[:2] - stitch_buf = canvas_cache.get("stitch_buf") - if stitch_buf is None or stitch_buf.shape[0] != rh or stitch_buf.shape[1] != rw: - stitch_buf = np.empty((rh, rw, 3), dtype=np.uint8) - canvas_cache["stitch_buf"] = stitch_buf - cv2.convertScaleAbs(roi, dst=stitch_buf) - stitched = stitch_buf - t_after_finalize = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 - - if do_profile: - xmf = canvas_cache.get("x_min_f", 0.0) - ymf = canvas_cache.get("y_min_f", 0.0) - print( - f"stitch_feature_optimized: canvas out_w={out_w}, out_h={out_h}, " - f"x_min={xmf:.2f}, y_min={ymf:.2f}" - ) - print( - "stitch_feature_optimized steps: " - f"brightness={t_after_brightness - t_profile0:.6f}s, " - f"geom={t_after_geom - t_after_brightness:.6f}s ({'hit' if geom_cache_hit else 'miss'}), " - f"reset={t_after_reset - t_after_geom:.6f}s, " - f"blend={t_after_blend - t_after_reset:.6f}s ({blended_tiles + 1} tiles incl ref), " - f"finalize={t_after_finalize - t_after_blend:.6f}s, " - f"total={t_after_finalize - t_profile0:.6f}s" - ) - - return stitched - - def find_stitch_boxes(self,img_src,show = False): - - img_corr = self.__correct_edge_bloom(img_src) - self.__detect_nonblack_regions(img_corr,show = show, row_tol = self.row_tol) - return len(self.boxes), self.boxes - - # return self.boxes - - - def _copy_into_tile_buffer(self, src_view, dst, flip_x=True, flip_y=True): - if flip_x and flip_y: - cv2.flip(src_view, -1, dst) - elif flip_x: - cv2.flip(src_view, 1, dst) - elif flip_y: - cv2.flip(src_view, 0, dst) - else: - np.copyto(dst, src_view) - - def get_stitch_unit(self, img, flip_x=True, flip_y=True,index = np.arange(15)): - """ - Extract each tile given by boxes (x0,y0,x1,y1), default flip both X and Y due to inverting lens. - - img: source image (use corrected image if you prefer) - - boxes: list of (x0,y0,x1,y1) - - flip_x, flip_y: booleans (performed sequentially: X then Y) - - overlay_original: if True paste onto original image copy, otherwise onto blank background - - bg_color: background color for blank canvas (B,G,R) - Returns: canvas image and list of flipped tiles as (x0,y0,flipped_tile) - """ - h, w = img.shape[:2] - active_index = index if self.units is None else self.index - if isinstance(active_index, int): - selected_index = {active_index} - else: - selected_index = set(active_index) - flipped_units = [] if self.units is None else self.units - - for i, (x0, y0, x1, y1) in enumerate(self.boxes): - if i not in selected_index: - continue - - x0 = max(0, int(x0)); y0 = max(0, int(y0)) - x1 = min(w, int(x1)); y1 = min(h, int(y1)) - if x1 <= x0 or y1 <= y0: - continue - - src_view = img[y0:y1, x0:x1] - if self.units is None: - tile = np.empty_like(src_view) - self._copy_into_tile_buffer(src_view, tile, flip_x=flip_x, flip_y=flip_y) - flipped_units.append((x0, y0, tile)) - else: - _, _, tile = flipped_units[i] - if tile.shape != src_view.shape or tile.dtype != src_view.dtype: - tile = np.empty_like(src_view) - self._copy_into_tile_buffer(src_view, tile, flip_x=flip_x, flip_y=flip_y) - flipped_units[i] = (x0, y0, tile) - - self.units = flipped_units - - return True - - def extract_tiles(self): - if not self.isstitch: - # tile_idx = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1]) - tile_idx = 7 - # if tile_idx is None: - # tile_idx = 7 - # Ensure index is always iterable (list or array) - self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx - self.ncol = 1 - self.nrow = 1 - - tile_list = self.__get_tile_by_index(self.index) - self.tile_list = tile_list - - return tile_list - - # def stitch_cal(self): - - # # init - # tile_list = self.tile_list - # img_stitched_height = None - # stitch_calibration = [] - # # Flip col and row for porting - # nrow, ncol = self.ncol, self.nrow - # if not self.isstitch: - # return self.tile_list[0], None - - # for i in range(ncol): - # img_stitched_width = None - # for j in range(nrow): - # if j == 0: - # img_j0 = tile_list[i*ncol] - # img_stitched_width = img_j0 - # else: - # img_1 = tile_list[i*ncol + j-1] - # img_2 = tile_list[i*ncol + j] - - # ov_pos, ov1, ov2, c_corr = self.__get_ov(img_1.astype(np.float32), img_2.astype(np.float32), dim=1) - - # stitch_calibration.append(["hor {}".format(j-1),ov_pos]) - - # # stitch images with stacking - # img_stitched_width = self.__stitch_optimized(img_stitched_width, img_2, ov_pos, ov1, ov2, blend=True, dim=1) - - # if i == 0: - # img_stitched_height = img_stitched_width - # stitch_calibration.append(["ver {}".format(i),0]) - # pass - # else: - # width = img_stitched_height.shape[1] - # width_new = img_stitched_width.shape[1] - # wdiff = width - width_new - # if wdiff > 0: - # img_stitched_width = np.hstack((img_stitched_width, img_stitched_width[:, -wdiff:])) - # elif wdiff < 0: - # img_stitched_width = img_stitched_width[:, :wdiff] - - # ov_pos_h, ov_h1, ov_h2, c_corr = self.__get_ov(img_stitched_height, img_stitched_width, dim=0) - - # stitch_calibration.append(["ver {}".format(i),ov_pos_h]) - - # img_stitched_height = self.__stitch_optimized(img_stitched_height, img_stitched_width, ov_pos_h, ov_h1, ov_h2, blend=True, dim=0) - - # self.stitch_calibration = stitch_calibration - - # return img_stitched_height, stitch_calibration - - # def __stitch_optimized(self, img1, img2, ov_pos, ov1, ov2, blend=True, dim=1): - # _dim = len(img1.shape) - - # # Clamp overlap for safety - # if dim == 1: - # ov_pos = int(max(0, min(ov_pos, img1.shape[1], img2.shape[1]))) - # else: - # ov_pos = int(max(0, min(ov_pos, img1.shape[0], img2.shape[0]))) - - # # Fast blend path with cached coefficient tensors - # if not blend or ov_pos <= 1: - # ov = ov1 - # else: - # blend_coef = self._get_blend_coef(ov_pos, dim, _dim) - # ov = ( - # ov1.astype(np.float32, copy=False) * blend_coef - # + ov2.astype(np.float32, copy=False) * (1.0 - blend_coef) - # ).astype(np.uint8) - - # # Build stitched output with one target allocation and slice writes - # if dim == 1: - # left_w = img1.shape[1] - ov_pos - # right_w = img2.shape[1] - ov_pos - # out_w = left_w + ov_pos + right_w - # if _dim == 3: - # stitched_img = np.empty((img1.shape[0], out_w, img1.shape[2]), dtype=img1.dtype) - # if left_w > 0: - # stitched_img[:, :left_w, :] = img1[:, :left_w, :] - # if ov_pos > 0: - # stitched_img[:, left_w:left_w + ov_pos, :] = ov - # if right_w > 0: - # stitched_img[:, left_w + ov_pos:, :] = img2[:, ov_pos:, :] - # else: - # stitched_img = np.empty((img1.shape[0], out_w), dtype=img1.dtype) - # if left_w > 0: - # stitched_img[:, :left_w] = img1[:, :left_w] - # if ov_pos > 0: - # stitched_img[:, left_w:left_w + ov_pos] = ov - # if right_w > 0: - # stitched_img[:, left_w + ov_pos:] = img2[:, ov_pos:] - # else: - # top_h = img1.shape[0] - ov_pos - # bottom_h = img2.shape[0] - ov_pos - # out_h = top_h + ov_pos + bottom_h - # if _dim == 3: - # stitched_img = np.empty((out_h, img1.shape[1], img1.shape[2]), dtype=img1.dtype) - # if top_h > 0: - # stitched_img[:top_h, :, :] = img1[:top_h, :, :] - # if ov_pos > 0: - # stitched_img[top_h:top_h + ov_pos, :, :] = ov - # if bottom_h > 0: - # stitched_img[top_h + ov_pos:, :, :] = img2[ov_pos:, :, :] - # else: - # stitched_img = np.empty((out_h, img1.shape[1]), dtype=img1.dtype) - # if top_h > 0: - # stitched_img[:top_h, :] = img1[:top_h, :] - # if ov_pos > 0: - # stitched_img[top_h:top_h + ov_pos, :] = ov - # if bottom_h > 0: - # stitched_img[top_h + ov_pos:, :] = img2[ov_pos:, :] - - # return stitched_img - - # def __get_ov(self, img1, img2, dim=1): - # ''' - # default dimension is along axis 1 (width of the images) - # ''' - # # pass parameters - # nominal_ov = self.nominal_ov_pos - # bound_ov = self.bound - - # cross_corr = [] - # h1, w1 = img1.shape[:2] - # h2, w2 = img2.shape[:2] - # # clamp tested overlap to valid range (1..min dimension) - # max_possible = w1 if dim == 1 else h1 - # max_possible = min(max_possible, w2 if dim == 1 else h2) - # lo = max(1, nominal_ov - bound_ov) - # hi = min(max_possible, nominal_ov + bound_ov) - # ov_space = range(lo, hi + 1) - - # for ov in ov_space: - # if dim == 1: - # ov_area1 = img1[:, max(0, w1 - ov):, :] - # ov_area2 = img2[:, :min(ov, w2), :] - # elif dim == 0: - # ov_area1 = img1[max(0, h1 - ov):, :, :] - # ov_area2 = img2[:min(ov, h2), :, :] - - # # require exact shape match (detect_nonblack_regions guarantees uniform tile size) - # if ov_area1.size == 0 or ov_area2.size == 0 or ov_area1.shape != ov_area2.shape: - # cross_corr.append(-1.0) - # continue - - # # multi-channel cross correlation (average over channels) - # corr = 0.0 - # for i in range(min(3, ov_area1.shape[2])): - # a = ov_area1[:, :, i].astype(np.float32) - # b = ov_area2[:, :, i].astype(np.float32) - # denom = np.sqrt(np.sum(a * a) * np.sum(b * b)) + 1e-12 - # corr += float(np.sum(a * b) / denom) - # corr = corr / float(min(3, ov_area1.shape[2])) - # cross_corr.append(corr) - - # # pick best overlap (ignore invalid entries) - # best_idx = int(np.argmax(cross_corr)) - # ov_pos = list(ov_space)[best_idx] - - # # extract the overlap regions (assume shapes match) - # if dim == 1: - # ov_1 = img1[:, max(0, w1 - ov_pos):, :].copy() - # ov_2 = img2[:, :min(ov_pos, w2), :].copy() - # else: - # ov_1 = img1[max(0, h1 - ov_pos):, :, :].copy() - # ov_2 = img2[:min(ov_pos, h2), :, :].copy() - # return ov_pos, ov_1, ov_2, cross_corr - - -# Helper functions for stitching - def __get_tile_index_by_position(self, x, y, tol=20): - """ - Find flipped tile whose top-left matches (x,y) within tol pixels. - flipped_units: list of (x,y,tile) - """ - i= 0 - for xi, yi, tile in self.units: - if abs(int(x) - int(xi)) <= tol and abs(int(y) - int(yi)) <= tol: - # return tile - return i - i += 1 - return None - - def __get_tile_by_index(self, idx): - """ - Retrieve tile by index - """ - if isinstance(idx, int): - idx = [idx] - - if len(idx) > len(self.units): - raise ValueError("Index exceeds number of tiles.") - - tile_list = [] - for i in range(len(idx)): - [xi, yi, tile] = self.units[idx[i]] - tile_list.append(tile) - - return tile_list - - def __correct_edge_bloom(self, img, method="divide", kernel_size=None, clip=True): - """ - Simple illumination correction to reduce edge bloom / vignetting. Only for getting tile boxes. - - method: "divide" (default) or "subtract" - * divide: img_corrected = (img / (bg + eps)) * gain -- preserves colors - * subtract: img_corrected = img - (bg - mean_bg) -- removes low-freq background - - kernel_size: size for large-kernel blur used as background estimate (odd int). - If None, chosen from image size. - Returns corrected uint8 image. - """ - if img is None: - raise ValueError("Input image is None.") - - h, w = img.shape[:2] - if kernel_size is None: - # choose large kernel relative to image size, ensure odd - k = max(31, min(h, w) // 2) - if k % 2 == 0: - k += 1 - kernel_size = k - # convert to float32 for stable ops - img_f = img.astype(np.float32) - # estimate background by large Gaussian blur (works well for smooth bloom) - bg = cv2.GaussianBlur(img_f, (kernel_size, kernel_size), 0) - if method == "subtract": - # remove low-frequency background but keep overall brightness - mean_bg = bg.mean(axis=(0, 1), keepdims=True) - corr = img_f - (bg - mean_bg) - else: - # divide: normalized gain correction - eps = 1e-6 - # compute gain so that mid-level stays near 128 - mean_bg = bg.mean(axis=(0, 1), keepdims=True) + eps - gain = (mean_bg * 0.5) # target scale; 0.5 keeps values reasonable - corr = (img_f / (bg + eps)) * gain - if clip: - corr = np.clip(corr, 0, 255).astype(np.uint8) - else: - corr = corr.astype(np.uint8) - return corr - - def __get_unit_pos(self,idx): - - if isinstance(idx, int): - idx = [idx] - - if len(idx) > len(self.units): - raise ValueError("Index exceeds number of vision units.") - - pos_list = [] - for i in range(len(idx)): - [xi, yi, tile] = self.units[idx[i]] - pos_list.append((xi, yi)) - - return pos_list - - # Computation Functions - def __detect_nonblack_regions(self, img, morph_k=3, n_boxes_target=15, bk_thresh_step=10, bk_thresh_max=100,show = False, row_tol = 0.4): - """ - Return binary mask of non-black regions and list of bounding boxes (x,y,w,h). - Uses adaptive thresholding on grayscale + small morphology to be robust to illumination. - - New behavior: after finding candidate boxes, center-crop all boxes to the same - width/height (smallest found) so returned boxes share a common size. - """ - if img is None: - raise ValueError("Input image is None.") - - h, w = img.shape[:2] - n_boxes = 0 - bk_thresh = 0 # initial black threshold - bk_thresh_step = 5 - bk_thresh_max = 150 - - # Also include a low global V threshold to capture colored but dark tiles - hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) - v = hsv[:,:,2] - ww = hh = 0 - while n_boxes != n_boxes_target and bk_thresh <= bk_thresh_max: - _, v_mask = cv2.threshold(v, bk_thresh, 255, cv2.THRESH_BINARY) # tune 50 if needed - - mask = v_mask.copy() - # Exclude left and right boundary (20 pixels each) from detection - edge = 20 - if w > edge * 2: - mask[:, :edge] = 0 - mask[:, -edge:] = 0 - - # Morphology: remove small specks, fill small holes - k = max(1, morph_k) - kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k, k)) - mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1) - mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) - # cv2.imshow("mask", mask) - # cv2.waitKey(0) - # find contours of candidate regions - contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - - boxes = [] - for c in contours: - x,y,ww,hh = cv2.boundingRect(c) - boxes.append((x,y,ww,hh)) - - # filter boxes by size (ignore too small or too large tiles) - boxes = self.__filter_boxes_by_size(boxes, (h, w)) - n_boxes = len(boxes) - bk_thresh += bk_thresh_step - - print("Detected {} candidate boxes with bk_thresh={}".format(n_boxes, bk_thresh - bk_thresh_step)) - if n_boxes == 0: - self.boxes = [] - return n_boxes - # compute common center-crop size: use minimum width/height among boxes - widths = [b[2] for b in boxes] - heights = [b[3] for b in boxes] - # target_w = max(1,int(min(widths))) - # target_h = max(1,int(min(heights))) - target_w = self.box_size[0] - target_h = self.box_size[1] - - # center-crop each box to (target_w, target_h) inside original box bounds - cropped_boxes = [] - for (x, y, ww, hh) in boxes: - cx = x + ww // 2 - cy = y + hh // 2 - nx0 = int(cx - target_w // 2) - ny0 = int(cy - target_h // 2) - # clamp to image bounds - nx0 = max(0, min(nx0, w - target_w)) - ny0 = max(0, min(ny0, h - target_h)) - # Convert to (x0, y0, x1, y1) format for direct image region retrieval - nx1 = nx0 + target_w - ny1 = ny0 + target_h - cropped_boxes.append((nx0, ny0, nx1, ny1)) - - # sort boxes left-to-right, top-to-bottom (useful for grid) - # Use quantized y-coordinates to group boxes into rows, making sorting robust to noise - if len(cropped_boxes) > 0: - # Calculate row tolerance based on typical box height (use median height) - heights = [b[3] - b[1] for b in cropped_boxes] - median_height = np.median(heights) if heights else target_h - # Use a fraction of the box height as tolerance (e.g., 20% of box height) - row_tolerance = max(1, int(median_height * row_tol)) - - # Quantize y-coordinates to group boxes into rows - def get_row_key(box): - y = box[1] # y0 coordinate - # Round y to nearest multiple of row_tolerance to group into rows - quantized_y = (y // row_tolerance) * row_tolerance - return (quantized_y, box[0]) # Sort by quantized row, then by x - - cropped_boxes = sorted(cropped_boxes, key=get_row_key) - else: - cropped_boxes = sorted(cropped_boxes, key=lambda b: (b[1], b[0])) - - self.boxes = cropped_boxes - img_with_boxes = self.draw_boxes(img= img) - if show: - cv2.imshow("Detected Boxes", img_with_boxes) - cv2.waitKey(0) - - img_path = "./config/compound_eye.png" - if os.path.exists(img_path): - cv2.imwrite(img_path, img_with_boxes) - else: - os.makedirs(os.path.dirname(img_path), exist_ok=True) - cv2.imwrite(img_path, img_with_boxes) - # print("Check if compound eye detection image valid: {}".format(img_path)) - return cropped_boxes - - def draw_boxes(self, img, thickness=2): - """ - Draw bounding boxes on the image. - - img: input image (will be copied, not modified in-place) - - thickness: line thickness - Returns: image with drawn boxes - """ - if self.boxes is None or len(self.boxes) == 0: - print("No boxes to draw") - return img - - img_draw = img.copy() - color=(0, 255, 0) - - for i, (x0, y0, x1, y1) in enumerate(self.boxes): - # Draw rectangle - cv2.rectangle(img_draw, (x0, y0), (x1, y1), color, thickness) - - # Draw index - text_size = cv2.getTextSize(str(i), cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0] - text_pos = (x0 + 5, y0 + text_size[1] + 5) - # Draw text background and text - cv2.rectangle(img_draw, (text_pos[0] - 2, text_pos[1] - text_size[1] - 2), - (text_pos[0] + text_size[0] + 2, text_pos[1] + 2), (0, 0, 0), -1) - cv2.putText(img_draw, str(i), text_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) - - return img_draw - - - - def __filter_boxes_by_size(self, boxes, img_shape): - """ - Remove boxes that are too small or too large relative to the image. - - boxes: list of (x,y,w,h) - - img_shape: (height, width) - - min_area_frac / max_area_frac: allowed area fraction of image - Returns filtered list of boxes. - """ - h, w = img_shape - img_area = h * w - min_area = img_area * float(self.min_area_frac) - max_area = img_area * float(self.max_area_frac) - filtered = [] - for (x,y,ww,hh) in boxes: - area = ww * hh - if area < min_area or area > max_area or (abs(ww - hh) > self.max_aspect_ratio): - continue - filtered.append((x,y,ww,hh)) - return filtered - - \ No newline at end of file +__all__ = ["Stitcher"] diff --git a/src/orisys/stitch/__init__.py b/src/orisys/stitch/__init__.py new file mode 100644 index 0000000..bcf579b --- /dev/null +++ b/src/orisys/stitch/__init__.py @@ -0,0 +1,33 @@ +"""Compound-eye stitching package (ops, backends, feature/box/tile modules).""" + +from .api import Stitcher +from .backends import available_blend_backends, get_blend_backend +from .config import BLEND_MODES, StitchStatus, resolve_blend_mode, validate_box_grid +from .ops import ( + build_blend_coef, + compute_final_crop_rect, + copy_into_tile_buffer, + coverage_ratio, + largest_fully_covered_rect, + match_bgr_brightness, + paste_tile_clipped, + tile_weight_map, +) + +__all__ = [ + "Stitcher", + "StitchStatus", + "BLEND_MODES", + "resolve_blend_mode", + "validate_box_grid", + "get_blend_backend", + "available_blend_backends", + "tile_weight_map", + "paste_tile_clipped", + "copy_into_tile_buffer", + "build_blend_coef", + "compute_final_crop_rect", + "largest_fully_covered_rect", + "match_bgr_brightness", + "coverage_ratio", +] diff --git a/src/orisys/stitch/api.py b/src/orisys/stitch/api.py new file mode 100644 index 0000000..0038dff --- /dev/null +++ b/src/orisys/stitch/api.py @@ -0,0 +1,324 @@ +"""Public Stitcher façade composing tiles, boxes, feature stitch, and blend backends.""" + +from __future__ import annotations + +import time +from typing import Optional + +import numpy as np + +from .backends import get_blend_backend +from .config import ( + BLEND_MODES, + StitchStatus, + numba_blend_available, + resolve_blend_mode, +) +from . import boxes as boxes_mod +from . import feature as feature_mod +from .ops import build_blend_coef +from . import tiles as tiles_mod + + +class Stitcher: + """Compound-eye stitcher (feature/ORB path + feathered blend backends).""" + + def __init__(self, isstitch=False, config_stitcher=None, verbose=False): + self.isstitch = isstitch + self.config_stitcher = config_stitcher + self.verbose = verbose + + self.min_area_frac = config_stitcher["min_area_frac"] + self.max_area_frac = config_stitcher["max_area_frac"] + self.max_aspect_ratio = config_stitcher["max_aspect_ratio"] + + self.stitch_valid = False + self.boxes = self.units = None + self.index = config_stitcher["index"] + self.stitch_calibration = self.tile_list = [] + self.ncol = config_stitcher["n"][0] + self.nrow = config_stitcher["n"][1] + self.nominal_ov_pos = config_stitcher["nominal_ov_pos"] + self.margin = config_stitcher["margin"] + self.bound = self.nominal_ov_pos - self.margin + self.row_tol = config_stitcher["row_tol"] + self.box_size = config_stitcher["box_size"] + # Physical tile count for box detection (often 15); grid n is the active stitch layout. + self.n_boxes_target = int(config_stitcher.get("n_boxes_target", 15)) + + self._blend_coef_cache = {} + self._feature_stitch_cache = { + "ref_idx": None, + "moving_indices": None, + "tile_shapes": None, + "transforms": None, + "canvas": {}, + } + self._feature_frame_idx = {"idx": 0} + self._feature_stitch_print_every = 30 + self._feature_stitch_pad = 5 + self._feature_stitch_match_brightness = False + self._blend_mode, fell_back = resolve_blend_mode(config_stitcher) + self._blend_backend_name = "numpy" if self._blend_mode == "none" else self._blend_mode + self._numba_failed = False + self._blend_backend = get_blend_backend( + self._blend_backend_name, + on_numba_failure=self._mark_numba_failed, + ) + self._last_coverage = 0.0 + self._last_ms = 0.0 + + if fell_back and self.isstitch: + print( + "Warning: stitcher blend=numba requested but numba is " + "unavailable; using numpy blend path." + ) + elif self.isstitch and self.verbose: + print(f"Stitch blend: {self._blend_mode}") + + if self.verbose: + print(f"Stitcher initialized with config: {self.config_stitcher}") + + @property + def blend_mode(self): + """Configured blend mode: ``none``, ``numpy``, or ``numba``.""" + return self._blend_mode + + @property + def status(self) -> StitchStatus: + return StitchStatus( + blend_mode=self._blend_mode, + backend_name=getattr(self._blend_backend, "name", self._blend_backend_name), + numba_failed=self._numba_failed, + coverage=float(self._last_coverage), + last_ms=float(self._last_ms), + tile_count=len(self.tile_list) if self.tile_list is not None else 0, + stitch_valid=bool(self.stitch_valid), + box_count=len(self.boxes) if self.boxes is not None else 0, + expected_boxes=int(self.n_boxes_target), + ) + + def configure_feature_stitch( + self, + *, + match_brightness=None, + pad=None, + print_every=None, + blend=None, + blend_backend=None, + ): + if match_brightness is not None: + self._feature_stitch_match_brightness = bool(match_brightness) + if pad is not None: + self._feature_stitch_pad = max(0, int(pad)) + if print_every is not None: + self._feature_stitch_print_every = max(1, int(print_every)) + mode = blend if blend is not None else blend_backend + if mode is not None: + mode = str(mode).strip().lower() + if mode not in BLEND_MODES: + raise ValueError( + f"Unknown blend={mode!r}; expected one of {sorted(BLEND_MODES)}" + ) + if mode == "numba" and (not numba_blend_available() or self._numba_failed): + print("Warning: numba unavailable; using numpy blend.") + mode = "numpy" + self._blend_mode = mode + self._blend_backend_name = "numpy" if mode == "none" else mode + if mode == "numba": + self._numba_failed = False + self._blend_backend = get_blend_backend( + self._blend_backend_name, + on_numba_failure=self._mark_numba_failed, + ) + + def _resolve_call_blend(self, blend): + if blend is None: + mode = self._blend_mode + elif isinstance(blend, bool): + if blend: + mode = self._blend_mode if self._blend_mode != "none" else "numpy" + else: + return False, self._blend_backend_name + else: + mode = str(blend).strip().lower() + if mode not in BLEND_MODES: + raise ValueError( + f"Unknown blend={mode!r}; expected one of {sorted(BLEND_MODES)}" + ) + if mode == "numba" and not numba_blend_available(): + mode = "numpy" + + if mode == "numba" and self._numba_failed: + mode = "numpy" + + if mode == "none": + return False, self._blend_backend_name + return True, mode + + def _mark_numba_failed(self, exc): + if self._numba_failed: + return + self._numba_failed = True + self._blend_backend_name = "numpy" + if self._blend_mode == "numba": + self._blend_mode = "numpy" + self._blend_backend = get_blend_backend("numpy") + print( + f"Warning: numba blend failed ({type(exc).__name__}: {exc}); " + f"falling back to numpy." + ) + + def _get_blend_coef(self, ov_pos, dim, ndim): + key = (ov_pos, dim, ndim) + coef = self._blend_coef_cache.get(key) + if coef is not None: + return coef + blend_coef = build_blend_coef(ov_pos, dim, ndim) + self._blend_coef_cache[key] = blend_coef + return blend_coef + + def _ensure_blend_backend(self, backend_name: str): + name = "numpy" if backend_name == "none" else backend_name + if self._numba_failed and name == "numba": + name = "numpy" + current = getattr(self._blend_backend, "name", None) + if current != name: + self._blend_backend = get_blend_backend( + name, + on_numba_failure=self._mark_numba_failed, + ) + self._blend_backend_name = name + + def stitch_rt_optimized(self, change_sv=False, s_mean=28.0, v_mean=180.0, blend=None): + """Stitch the current tile list via the feature/ORB path. + + Returns ``None`` when a valid feature layout cannot be built. + """ + use_feather, blend_backend = self._resolve_call_blend(blend) + prev_backend = self._blend_backend_name + self._ensure_blend_backend(blend_backend) + t0 = time.perf_counter() + try: + feature_stitched = self.stitch_feature_optimized( + blend=use_feather, + match_brightness=self._feature_stitch_match_brightness, + pad=self._feature_stitch_pad, + print_every=self._feature_stitch_print_every, + ) + finally: + if self._numba_failed: + self._ensure_blend_backend("numpy") + else: + self._ensure_blend_backend(prev_backend) + self._last_ms = (time.perf_counter() - t0) * 1000.0 + canvas = self._feature_stitch_cache.get("canvas") or {} + self._last_coverage = float(canvas.get("last_coverage", self._last_coverage)) + if feature_stitched is not None: + return feature_stitched + return None + + def stitch_feature_optimized(self, blend=True, match_brightness=False, pad=5, print_every=30): + return feature_mod.stitch_feature_optimized( + isstitch=self.isstitch, + index=self.index, + nrow=self.nrow, + ncol=self.ncol, + tile_list=self.tile_list, + cache=self._feature_stitch_cache, + blend_backend=self._blend_backend, + blend=blend, + match_brightness=match_brightness, + pad=pad, + print_every=print_every, + verbose=self.verbose, + frame_idx_holder=self._feature_frame_idx, + ) + + def find_stitch_boxes(self, img_src, show=False): + n_boxes, boxes, grid_ok = boxes_mod.find_stitch_boxes( + img_src, + box_size=self.box_size, + min_area_frac=self.min_area_frac, + max_area_frac=self.max_area_frac, + max_aspect_ratio=self.max_aspect_ratio, + row_tol=self.row_tol, + nrow=self.nrow, + ncol=self.ncol, + show=show, + n_boxes_target=self.n_boxes_target, + ) + self.boxes = boxes + # Detection target is physical unit count; active stitch grid is nrow*ncol via index. + self.stitch_valid = bool(self.isstitch and n_boxes == self.n_boxes_target) + if self.isstitch and n_boxes != self.n_boxes_target: + print( + f"Warning: stitch box count {n_boxes} != expected " + f"n_boxes_target={self.n_boxes_target}; stitch_valid=False." + ) + elif self.isstitch and not grid_ok and self.verbose: + # Soft note: active index length may be nrow*ncol while more boxes exist. + print( + f"Note: detected {n_boxes} boxes; active stitch grid is " + f"{self.nrow}x{self.ncol} via index ({len(tiles_mod.normalize_index(self.index))} tiles)." + ) + return n_boxes, self.boxes + + def get_stitch_unit(self, img, flip_x=True, flip_y=True, index=np.arange(15)): + active_index = index if self.units is None else self.index + if isinstance(active_index, int): + selected_index = {active_index} + else: + selected_index = set(int(i) for i in active_index) + self.units = tiles_mod.extract_units_from_boxes( + img, + self.boxes, + selected_index=selected_index, + units=self.units, + flip_x=flip_x, + flip_y=flip_y, + ) + return True + + def extract_tiles(self): + if not self.isstitch: + tile_idx = 7 + self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx + self.ncol = 1 + self.nrow = 1 + + tile_list = tiles_mod.get_tile_by_index(self.units, self.index) + self.tile_list = tile_list + return tile_list + + def draw_boxes(self, img, thickness=2): + return boxes_mod.draw_boxes(img, self.boxes, thickness=thickness) + + # Compatibility aliases used by older call sites / tests. + def _copy_into_tile_buffer(self, src_view, dst, flip_x=True, flip_y=True): + from .ops import copy_into_tile_buffer + + copy_into_tile_buffer(src_view, dst, flip_x=flip_x, flip_y=flip_y) + + def _tile_weight_map(self, h, w, margin=16): + from .ops import tile_weight_map + + return tile_weight_map(h, w, margin) + + def _paste_tile_clipped(self, canvas, tile, y0, x0, cov=None): + from .ops import paste_tile_clipped + + paste_tile_clipped(canvas, tile, y0, x0, cov=cov) + + def _blend_tile_inplace(self, blend, wsum, cov, src, weight, y0, x0, scratch=None): + self._blend_backend.blend_tile_inplace(blend, wsum, cov, src, weight, y0, x0) + + def _compute_final_crop_rect(self, cov): + from .ops import compute_final_crop_rect + + return compute_final_crop_rect(cov) + + def _largest_fully_covered_rect(self, cov): + from .ops import largest_fully_covered_rect + + return largest_fully_covered_rect(cov) diff --git a/src/orisys/stitch/backends/__init__.py b/src/orisys/stitch/backends/__init__.py new file mode 100644 index 0000000..c251c75 --- /dev/null +++ b/src/orisys/stitch/backends/__init__.py @@ -0,0 +1,47 @@ +"""Blend backend registry (numpy / numba; future: torch / cupy).""" + +from __future__ import annotations + +from typing import Callable, Optional + +from .base import BlendBackend +from .numba_backend import NumbaBlendBackend +from .numpy_backend import NumpyBlendBackend +from ..config import numba_blend_available + +_BACKENDS = { + "numpy": NumpyBlendBackend, + "numba": NumbaBlendBackend, +} + + +def get_blend_backend( + name: str, + *, + on_numba_failure: Optional[Callable[[BaseException], None]] = None, +) -> BlendBackend: + """Return a blend backend instance by name. + + Unknown names raise ``ValueError``. ``numba`` falls back to construction + even if JIT is missing; the backend itself falls back to numpy on failure. + """ + key = str(name).strip().lower() + if key == "none": + # Feather off uses paste path; keep a numpy backend available if forced. + key = "numpy" + if key not in _BACKENDS: + raise ValueError( + f"Unknown blend backend {name!r}; expected one of " + f"{sorted(set(_BACKENDS) | {'none'})} " + f"(future backends e.g. 'torch' / 'cupy' are not registered yet)." + ) + if key == "numba": + return NumbaBlendBackend(on_failure=on_numba_failure) + return NumpyBlendBackend() + + +def available_blend_backends() -> list[str]: + names = ["none", "numpy"] + if numba_blend_available(): + names.append("numba") + return names diff --git a/src/orisys/stitch/backends/base.py b/src/orisys/stitch/backends/base.py new file mode 100644 index 0000000..4b6fa97 --- /dev/null +++ b/src/orisys/stitch/backends/base.py @@ -0,0 +1,29 @@ +"""BlendBackend protocol for CPU (and future GPU) stitch compositors.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class BlendBackend(Protocol): + """Array-in / in-place composite kernel. + + Implementations must accept contiguous float32/uint8 NumPy arrays today. + Future Torch/CuPy backends can adapt tensors to this contract or extend it. + """ + + name: str + + def blend_tile_inplace( + self, + blend, + wsum, + cov, + src, + weight, + y0: int, + x0: int, + ) -> None: + """Composite ``src`` weighted by ``weight`` onto canvas buffers at (y0, x0).""" + ... diff --git a/src/orisys/stitch/backends/numba_backend.py b/src/orisys/stitch/backends/numba_backend.py new file mode 100644 index 0000000..17a73af --- /dev/null +++ b/src/orisys/stitch/backends/numba_backend.py @@ -0,0 +1,53 @@ +"""Numba blend backend wrapping ``orisys_blend.kernel``.""" + +from __future__ import annotations + +import numpy as np + +from .numpy_backend import NumpyBlendBackend, blend_kernel_numpy + +try: + from orisys_blend.kernel import _NUMBA_AVAILABLE, blend_kernel as _blend_kernel_numba +except Exception: + _NUMBA_AVAILABLE = False + _blend_kernel_numba = None + +if _blend_kernel_numba is None: + _NUMBA_AVAILABLE = False + + +class NumbaBlendBackend: + name = "numba" + + def __init__(self, *, on_failure=None): + self._on_failure = on_failure + self._fallback = NumpyBlendBackend() + self.failed = False + + def blend_tile_inplace(self, blend, wsum, cov, src, weight, y0, x0) -> None: + h, w = src.shape[:2] + H, W = blend.shape[:2] + y0d, x0d = max(0, y0), max(0, x0) + y1d, x1d = min(H, y0 + h), min(W, x0 + w) + if y0d >= y1d or x0d >= x1d: + return + + sy0, sx0 = y0d - y0, x0d - x0 + sh, sw = y1d - y0d, x1d - x0d + src_f = src[sy0 : sy0 + sh, sx0 : sx0 + sw].astype(np.float32, copy=False) + args = ( + blend[y0d:y1d, x0d:x1d], + wsum[y0d:y1d, x0d:x1d], + cov[y0d:y1d, x0d:x1d], + src_f, + weight[sy0 : sy0 + sh, sx0 : sx0 + sw], + ) + if not self.failed and _NUMBA_AVAILABLE and _blend_kernel_numba is not None: + try: + _blend_kernel_numba(*args) + return + except Exception as exc: + self.failed = True + if self._on_failure is not None: + self._on_failure(exc) + blend_kernel_numpy(*args) diff --git a/src/orisys/stitch/backends/numpy_backend.py b/src/orisys/stitch/backends/numpy_backend.py new file mode 100644 index 0000000..c76f66a --- /dev/null +++ b/src/orisys/stitch/backends/numpy_backend.py @@ -0,0 +1,61 @@ +"""NumPy blend backend (no JIT).""" + +from __future__ import annotations + +import numpy as np + + +def blend_kernel_numpy(blend_patch, wsum_patch, cov_patch, src_patch, w_patch): + """Incremental weighted-average composite (pure NumPy, no JIT). + + Uses contiguous arithmetic instead of boolean fancy indexing. Callers must + zero ``blend`` / ``wsum`` on uncovered pixels (as the stitcher reset does); + then the covered lerp also paints first-touch pixels correctly when ``w > 0``. + """ + w = w_patch if w_patch.dtype == np.float32 else np.asarray(w_patch, dtype=np.float32) + src = src_patch if src_patch.dtype == np.float32 else np.asarray(src_patch, dtype=np.float32) + + uncovered = cov_patch == 0 + if uncovered.all(): + np.copyto(blend_patch, src) + np.copyto(wsum_patch, w) + cov_patch.fill(1) + return + + ws = wsum_patch + w + a = w / np.maximum(ws, np.float32(1e-6)) + for c in range(src.shape[2]): + bc = blend_patch[..., c] + bc += (src[..., c] - bc) * a + np.copyto(wsum_patch, ws) + + if uncovered.any(): + need_fix = uncovered & (w < np.float32(1e-6)) + if need_fix.any(): + blend_patch[need_fix] = src[need_fix] + wsum_patch[need_fix] = w[need_fix] + + cov_patch.fill(1) + + +class NumpyBlendBackend: + name = "numpy" + + def blend_tile_inplace(self, blend, wsum, cov, src, weight, y0, x0) -> None: + h, w = src.shape[:2] + H, W = blend.shape[:2] + y0d, x0d = max(0, y0), max(0, x0) + y1d, x1d = min(H, y0 + h), min(W, x0 + w) + if y0d >= y1d or x0d >= x1d: + return + + sy0, sx0 = y0d - y0, x0d - x0 + sh, sw = y1d - y0d, x1d - x0d + src_f = src[sy0 : sy0 + sh, sx0 : sx0 + sw].astype(np.float32, copy=False) + blend_kernel_numpy( + blend[y0d:y1d, x0d:x1d], + wsum[y0d:y1d, x0d:x1d], + cov[y0d:y1d, x0d:x1d], + src_f, + weight[sy0 : sy0 + sh, sx0 : sx0 + sw], + ) diff --git a/src/orisys/stitch/boxes.py b/src/orisys/stitch/boxes.py new file mode 100644 index 0000000..c69a4a6 --- /dev/null +++ b/src/orisys/stitch/boxes.py @@ -0,0 +1,217 @@ +"""Box discovery and visualization for compound-eye calibration.""" + +from __future__ import annotations + +import os +from typing import List, Optional, Sequence, Tuple + +import cv2 +import numpy as np + + +def correct_edge_bloom(img, method: str = "divide", kernel_size=None, clip: bool = True): + """Illumination correction to reduce edge bloom (box detection only).""" + if img is None: + raise ValueError("Input image is None.") + + h, w = img.shape[:2] + if kernel_size is None: + k = max(31, min(h, w) // 2) + if k % 2 == 0: + k += 1 + kernel_size = k + img_f = img.astype(np.float32) + bg = cv2.GaussianBlur(img_f, (kernel_size, kernel_size), 0) + if method == "subtract": + mean_bg = bg.mean(axis=(0, 1), keepdims=True) + corr = img_f - (bg - mean_bg) + else: + eps = 1e-6 + mean_bg = bg.mean(axis=(0, 1), keepdims=True) + eps + gain = mean_bg * 0.5 + corr = (img_f / (bg + eps)) * gain + if clip: + corr = np.clip(corr, 0, 255).astype(np.uint8) + else: + corr = corr.astype(np.uint8) + return corr + + +def filter_boxes_by_size( + boxes: Sequence[Tuple[int, int, int, int]], + img_shape, + *, + min_area_frac: float, + max_area_frac: float, + max_aspect_ratio: float, +) -> List[Tuple[int, int, int, int]]: + h, w = img_shape + img_area = h * w + min_area = img_area * float(min_area_frac) + max_area = img_area * float(max_area_frac) + filtered = [] + for (x, y, ww, hh) in boxes: + area = ww * hh + if area < min_area or area > max_area or (abs(ww - hh) > max_aspect_ratio): + continue + filtered.append((x, y, ww, hh)) + return filtered + + +def detect_nonblack_regions( + img, + *, + box_size: Sequence[int], + min_area_frac: float, + max_area_frac: float, + max_aspect_ratio: float, + morph_k: int = 3, + n_boxes_target: int = 15, + bk_thresh_step: int = 5, + bk_thresh_max: int = 150, + show: bool = False, + row_tol: float = 0.4, + save_preview_path: str = "./config/compound_eye.png", +) -> Tuple[int, List[Tuple[int, int, int, int]]]: + """Detect tile boxes; returns ``(n_boxes, boxes as x0,y0,x1,y1)``.""" + if img is None: + raise ValueError("Input image is None.") + + h, w = img.shape[:2] + n_boxes = 0 + bk_thresh = 0 + boxes: List[Tuple[int, int, int, int]] = [] + + hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) + v = hsv[:, :, 2] + while n_boxes != n_boxes_target and bk_thresh <= bk_thresh_max: + _, v_mask = cv2.threshold(v, bk_thresh, 255, cv2.THRESH_BINARY) + mask = v_mask.copy() + edge = 20 + if w > edge * 2: + mask[:, :edge] = 0 + mask[:, -edge:] = 0 + + k = max(1, morph_k) + kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k, k)) + mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1) + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) + contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + boxes = [] + for c in contours: + x, y, ww, hh = cv2.boundingRect(c) + boxes.append((x, y, ww, hh)) + + boxes = filter_boxes_by_size( + boxes, + (h, w), + min_area_frac=min_area_frac, + max_area_frac=max_area_frac, + max_aspect_ratio=max_aspect_ratio, + ) + n_boxes = len(boxes) + bk_thresh += bk_thresh_step + + print( + "Detected {} candidate boxes with bk_thresh={}".format( + n_boxes, bk_thresh - bk_thresh_step + ) + ) + if n_boxes == 0: + return 0, [] + + target_w = int(box_size[0]) + target_h = int(box_size[1]) + cropped_boxes = [] + for (x, y, ww, hh) in boxes: + cx = x + ww // 2 + cy = y + hh // 2 + nx0 = int(cx - target_w // 2) + ny0 = int(cy - target_h // 2) + nx0 = max(0, min(nx0, w - target_w)) + ny0 = max(0, min(ny0, h - target_h)) + cropped_boxes.append((nx0, ny0, nx0 + target_w, ny0 + target_h)) + + if cropped_boxes: + heights = [b[3] - b[1] for b in cropped_boxes] + median_height = np.median(heights) if heights else target_h + row_tolerance = max(1, int(median_height * row_tol)) + + def get_row_key(box): + y = box[1] + quantized_y = (y // row_tolerance) * row_tolerance + return (quantized_y, box[0]) + + cropped_boxes = sorted(cropped_boxes, key=get_row_key) + else: + cropped_boxes = sorted(cropped_boxes, key=lambda b: (b[1], b[0])) + + img_with_boxes = draw_boxes(img, cropped_boxes) + if show: + cv2.imshow("Detected Boxes", img_with_boxes) + cv2.waitKey(0) + + if save_preview_path: + parent = os.path.dirname(save_preview_path) + if parent: + os.makedirs(parent, exist_ok=True) + cv2.imwrite(save_preview_path, img_with_boxes) + + return len(cropped_boxes), cropped_boxes + + +def draw_boxes(img, boxes, thickness: int = 2): + if boxes is None or len(boxes) == 0: + print("No boxes to draw") + return img + + img_draw = img.copy() + color = (0, 255, 0) + for i, (x0, y0, x1, y1) in enumerate(boxes): + cv2.rectangle(img_draw, (x0, y0), (x1, y1), color, thickness) + text_size = cv2.getTextSize(str(i), cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0] + text_pos = (x0 + 5, y0 + text_size[1] + 5) + cv2.rectangle( + img_draw, + (text_pos[0] - 2, text_pos[1] - text_size[1] - 2), + (text_pos[0] + text_size[0] + 2, text_pos[1] + 2), + (0, 0, 0), + -1, + ) + cv2.putText(img_draw, str(i), text_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) + return img_draw + + +def find_stitch_boxes( + img_src, + *, + box_size, + min_area_frac, + max_area_frac, + max_aspect_ratio, + row_tol, + nrow: int, + ncol: int, + show: bool = False, + n_boxes_target: Optional[int] = None, +): + """Run bloom correction + detection; return ``(n_boxes, boxes, grid_ok)``. + + ``n_boxes_target`` is the physical tile count to detect (default 15). + ``grid_ok`` is True when that count matches ``n_boxes_target``. + """ + img_corr = correct_edge_bloom(img_src) + target = int(n_boxes_target) if n_boxes_target is not None else 15 + n_boxes, boxes = detect_nonblack_regions( + img_corr, + box_size=box_size, + min_area_frac=min_area_frac, + max_area_frac=max_area_frac, + max_aspect_ratio=max_aspect_ratio, + n_boxes_target=target, + show=show, + row_tol=row_tol, + ) + grid_ok = n_boxes == target + return n_boxes, boxes, grid_ok diff --git a/src/orisys/stitch/config.py b/src/orisys/stitch/config.py new file mode 100644 index 0000000..77cbe0d --- /dev/null +++ b/src/orisys/stitch/config.py @@ -0,0 +1,76 @@ +"""Stitch configuration, blend-mode resolution, and status.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any, Mapping, Optional + +# JSON ``stitcher.blend`` / ``blend_list``: none = fast paste; numpy/numba = feathered. +BLEND_MODES = frozenset({"none", "numpy", "numba"}) +DEFAULT_BLEND_MODE = "none" + + +def numba_blend_available() -> bool: + try: + from orisys_blend.kernel import _NUMBA_AVAILABLE, blend_kernel + except Exception: + return False + return bool(_NUMBA_AVAILABLE and blend_kernel is not None) + + +def resolve_blend_mode(config_stitcher: Optional[Mapping[str, Any]] = None): + """Resolve stitcher blend mode from env, then JSON ``blend``. + + Returns ``(mode, fell_back_from_numba)`` where mode is one of + ``none`` / ``numpy`` / ``numba``. + + Priority: ``ORISYS_STITCH_BLEND``, ``ORISYS_STITCH_BLEND_BACKEND``, + then ``config_stitcher["blend"]`` / legacy ``blend_backend``, then ``none``. + """ + requested = None + for key in ("ORISYS_STITCH_BLEND", "ORISYS_STITCH_BLEND_BACKEND"): + env = os.environ.get(key, "").strip().lower() + if env in BLEND_MODES: + requested = env + break + + if requested is None and isinstance(config_stitcher, dict): + for key in ("blend", "blend_backend"): + cfg = str(config_stitcher.get(key, "")).strip().lower() + if cfg in BLEND_MODES: + requested = cfg + break + + if requested is None: + requested = DEFAULT_BLEND_MODE + + if requested == "numba" and not numba_blend_available(): + return "numpy", True + return requested, False + + +def expected_box_count(nrow: int, ncol: int) -> int: + return int(nrow) * int(ncol) + + +def validate_box_grid(boxes, nrow: int, ncol: int) -> bool: + """True when detected boxes match the configured grid size.""" + if boxes is None: + return False + return len(boxes) == expected_box_count(nrow, ncol) + + +@dataclass +class StitchStatus: + """Runtime health snapshot for compound-eye stitching.""" + + blend_mode: str = DEFAULT_BLEND_MODE + backend_name: str = "numpy" + numba_failed: bool = False + coverage: float = 0.0 + last_ms: float = 0.0 + tile_count: int = 0 + stitch_valid: bool = False + box_count: int = 0 + expected_boxes: int = 0 diff --git a/src/orisys/stitch/feature.py b/src/orisys/stitch/feature.py new file mode 100644 index 0000000..a754fe8 --- /dev/null +++ b/src/orisys/stitch/feature.py @@ -0,0 +1,599 @@ +"""Feature-based (ORB/BFMatcher) stitch plan and render orchestration.""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Sequence + +import cv2 +import numpy as np + +from .backends.base import BlendBackend +from .ops import ( + compute_final_crop_rect, + coverage_ratio, + match_bgr_brightness, + mean_gray_bgr, + paste_tile_clipped, + tile_weight_map, + translation_from_H_to_ref, +) +from .tiles import normalize_index + + +def tile_to_bgr(img): + if img.ndim == 2: + return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + return img + + +def tile_to_gray(bgr): + if bgr.ndim == 2: + return bgr + return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) + + +def ensure_u8_gray(gray): + if gray.dtype == np.uint8: + return gray + return cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) + + +def preprocess_gray_for_orb(g0, g1): + g0 = cv2.GaussianBlur(g0, (5, 5), 0) + g1 = cv2.GaussianBlur(g1, (5, 5), 0) + g0 = cv2.equalizeHist(g0) + g1 = cv2.equalizeHist(g1) + return g0, g1 + + +def estimate_bfmatcher_stitch_transform( + t0, + t1, + *, + ransac_reproj_thresh=5.0, + ratio=0.9, + orb_nfeatures=5000, + verbose=True, +): + """Estimate a 3x3 transform mapping tile coordinates into the reference frame.""" + t0_bgr = tile_to_bgr(t0) + t1_bgr = tile_to_bgr(t1) + g0 = ensure_u8_gray(tile_to_gray(t0_bgr)) + g1 = ensure_u8_gray(tile_to_gray(t1_bgr)) + + if verbose: + print( + f"BF stitcher debug: t0 dtype={t0.dtype}, t1 dtype={t1.dtype}, " + f"g0 dtype={g0.dtype}, g0 range=({np.min(g0)}, {np.max(g0)}), " + f"g1 dtype={g1.dtype}, g1 range=({np.min(g1)}, {np.max(g1)})" + ) + + # g0, g1 = preprocess_gray_for_orb(g0, g1) + orb = cv2.ORB_create( + nfeatures=orb_nfeatures, + scaleFactor=1.2, + nlevels=8, + edgeThreshold=15, + fastThreshold=5, + ) + k0, d0 = orb.detectAndCompute(g0, None) + k1, d1 = orb.detectAndCompute(g1, None) + + info = { + "ok": False, + "method": None, + "inliers": 0, + "mean_err": None, + "good_matches": 0, + "k0": k0, + "k1": k1, + "t0_bgr": t0_bgr, + "t1_bgr": t1_bgr, + } + if d0 is None or d1 is None or len(k0) < 4 or len(k1) < 4: + if verbose: + print( + "OpenCV BFMatcher stitcher: not enough features.", + f"k0={0 if k0 is None else len(k0)}", + f"k1={0 if k1 is None else len(k1)}", + f"d0_is_none={d0 is None}", + f"d1_is_none={d1 is None}", + ) + return None, info + + matches = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False).knnMatch(d0, d1, k=2) + good = [] + for m_n in matches: + if len(m_n) != 2: + continue + m, n = m_n + if m.distance < ratio * n.distance: + good.append(m) + + info["good_matches"] = len(good) + if verbose: + print(f"BFMatcher good matches: {len(good)}") + if len(good) < 4: + return None, info + + pts0 = np.float32([k0[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) + pts1 = np.float32([k1[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) + method = "affine_partial2d" + M, inlier_mask = cv2.estimateAffinePartial2D( + pts1, + pts0, + method=cv2.RANSAC, + ransacReprojThreshold=ransac_reproj_thresh, + ) + if M is not None: + H = np.vstack([M, [0, 0, 1]]).astype(np.float64) + else: + method = "homography" + H, inlier_mask = cv2.findHomography(pts1, pts0, cv2.RANSAC, ransac_reproj_thresh) + + info["method"] = method + info["inliers"] = int(inlier_mask.sum()) if inlier_mask is not None else 0 + if verbose: + print(f"Estimated {method} inliers: {info['inliers']}") + if H is None: + return None, info + + if abs(H[2, 2]) > 1e-12: + H = H / H[2, 2] + + mean_err = 1e9 + if inlier_mask is not None: + inlier_bool = inlier_mask.ravel().astype(bool) + if np.any(inlier_bool): + projected = cv2.perspectiveTransform(pts1[inlier_bool], H) + err = np.linalg.norm( + projected.reshape(-1, 2) - pts0[inlier_bool].reshape(-1, 2), + axis=1, + ) + mean_err = float(err.mean()) if err.size else 1e9 + info["mean_err"] = mean_err + info["ok"] = True + return H, info + + +def get_feature_stitch_plan(isstitch, index, nrow, ncol, tile_list): + if not isstitch: + return None + + configured_indices = normalize_index(index) + total_tiles = len(configured_indices) + expected_tiles = nrow * ncol + if total_tiles <= 1 or total_tiles != expected_tiles or len(tile_list) < total_tiles: + return None + + ref_pos = (nrow // 2) * ncol + (ncol // 2) + if ref_pos >= total_tiles: + return None + + selected_positions = tuple(range(total_tiles)) + ref_idx = configured_indices[ref_pos] + moving_positions = tuple(pos for pos in selected_positions if pos != ref_pos) + moving_indices = tuple(configured_indices[pos] for pos in moving_positions) + moving_tiles = [tile_list[pos] for pos in moving_positions] + return ref_idx, tile_list[ref_pos], moving_indices, moving_tiles, selected_positions + + +def _geom_cache_key(t_ref_bgr, tiles_bgr_c, pad): + ref_h, ref_w = t_ref_bgr.shape[:2] + shapes = tuple(tb.shape[:2] for tb in tiles_bgr_c) + return ref_h, ref_w, shapes, pad + + +def ensure_feature_geom_cache( + canvas_cache, + t_ref_bgr, + tiles_bgr_c, + H_use_list, + pad, + *, + allocate_blend_bufs=True, +): + h0c, w0c = t_ref_bgr.shape[:2] + geom_key = _geom_cache_key(t_ref_bgr, tiles_bgr_c, pad) + geom_cache_hit = canvas_cache.get("geom_key") == geom_key + if geom_cache_hit: + if allocate_blend_bufs and ( + "blend_buf" not in canvas_cache + or "wsum_buf" not in canvas_cache + or "cov_buf" not in canvas_cache + ): + out_h = canvas_cache["out_h"] + out_w = canvas_cache["out_w"] + canvas_cache["blend_buf"] = np.zeros((out_h, out_w, 3), dtype=np.float32) + canvas_cache["wsum_buf"] = np.zeros((out_h, out_w), dtype=np.float32) + canvas_cache["cov_buf"] = np.zeros((out_h, out_w), dtype=np.uint8) + return True + + ref_corners = np.array([[0, 0], [w0c, 0], [w0c, h0c], [0, h0c]], dtype=np.float64) + pts = [ref_corners] + for tile_bgr, H_use in zip(tiles_bgr_c, H_use_list): + if H_use is None: + continue + hi, wi = tile_bgr.shape[:2] + tx, ty = translation_from_H_to_ref(H_use) + pts.append( + np.array( + [[tx, ty], [tx + wi, ty], [tx + wi, ty + hi], [tx, ty + hi]], + dtype=np.float64, + ) + ) + + all_xy = np.concatenate(pts, axis=0) + x_min_f = float(all_xy[:, 0].min()) - pad + y_min_f = float(all_xy[:, 1].min()) - pad + x_max_f = float(all_xy[:, 0].max()) + pad + y_max_f = float(all_xy[:, 1].max()) + pad + out_w = max(1, int(np.ceil(x_max_f - x_min_f))) + out_h = max(1, int(np.ceil(y_max_f - y_min_f))) + ref_y0 = int(np.floor(-y_min_f)) + ref_x0 = int(np.floor(-x_min_f)) + + pastes = [] + for H_use in H_use_list: + if H_use is None: + pastes.append(None) + continue + tx, ty = translation_from_H_to_ref(H_use) + pastes.append((int(np.floor(ty - y_min_f)), int(np.floor(tx - x_min_f)))) + + _dr_y0 = max(0, ref_y0) + _dr_x0 = max(0, ref_x0) + _dr_y1 = min(out_h, ref_y0 + h0c) + _dr_x1 = min(out_w, ref_x0 + w0c) + for _tb, _pa in zip(tiles_bgr_c, pastes): + if _pa is None: + continue + _py, _px = _pa + _th, _tw = _tb.shape[:2] + _dr_y0 = min(_dr_y0, max(0, _py)) + _dr_x0 = min(_dr_x0, max(0, _px)) + _dr_y1 = max(_dr_y1, min(out_h, _py + _th)) + _dr_x1 = max(_dr_x1, min(out_w, _px + _tw)) + + canvas_cache["geom_key"] = geom_key + canvas_cache["out_w"] = out_w + canvas_cache["out_h"] = out_h + canvas_cache["ref_y0"] = ref_y0 + canvas_cache["ref_x0"] = ref_x0 + canvas_cache["pastes"] = pastes + canvas_cache["x_min_f"] = x_min_f + canvas_cache["y_min_f"] = y_min_f + canvas_cache["crop_rect"] = None + canvas_cache["dirty_rect"] = (_dr_y0, _dr_y1, _dr_x0, _dr_x1) + canvas_cache["weight_maps"] = {} + canvas_cache["uint8_buf"] = np.zeros((out_h, out_w, 3), dtype=np.uint8) + if allocate_blend_bufs: + canvas_cache["blend_buf"] = np.zeros((out_h, out_w, 3), dtype=np.float32) + canvas_cache["wsum_buf"] = np.zeros((out_h, out_w), dtype=np.float32) + canvas_cache["cov_buf"] = np.zeros((out_h, out_w), dtype=np.uint8) + else: + canvas_cache["cov_buf"] = np.zeros((out_h, out_w), dtype=np.uint8) + canvas_cache.pop("blend_buf", None) + canvas_cache.pop("wsum_buf", None) + return False + + +def paste_feature_tiles_fast( + t_ref_bgr, + tiles_bgr_c, + H_use_list, + canvas_cache, + *, + geom_cache_hit, + frame_idx=0, + print_every=30, + verbose=False, + t_profile0=0.0, + t_after_brightness=0.0, + t_after_geom=0.0, +): + do_profile = verbose and frame_idx % print_every == 0 + out_w = canvas_cache["out_w"] + out_h = canvas_cache["out_h"] + ref_y0 = canvas_cache["ref_y0"] + ref_x0 = canvas_cache["ref_x0"] + pastes = canvas_cache["pastes"] + + canvas = canvas_cache.get("uint8_buf") + if canvas is None or canvas.shape[0] != out_h or canvas.shape[1] != out_w: + canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8) + canvas_cache["uint8_buf"] = canvas + + dirty = canvas_cache.get("dirty_rect") + if geom_cache_hit and dirty is not None: + dy0, dy1, dx0, dx1 = dirty + canvas[dy0:dy1, dx0:dx1] = 0 + else: + canvas.fill(0) + t_after_reset = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + need_crop = canvas_cache.get("crop_rect") is None + cov = canvas_cache.get("cov_buf") if need_crop else None + if need_crop: + if cov is None or cov.shape != (out_h, out_w): + cov = np.zeros((out_h, out_w), dtype=np.uint8) + canvas_cache["cov_buf"] = cov + else: + if dirty is not None: + dy0, dy1, dx0, dx1 = dirty + cov[dy0:dy1, dx0:dx1] = 0 + else: + cov.fill(0) + + pasted_tiles = 0 + for tile_bgr, H_use, paste in zip(tiles_bgr_c, H_use_list, pastes): + if H_use is None or paste is None: + continue + y0, x0 = paste + paste_tile_clipped(canvas, tile_bgr, y0, x0, cov=cov) + pasted_tiles += 1 + paste_tile_clipped(canvas, t_ref_bgr, ref_y0, ref_x0, cov=cov) + t_after_paste = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + crop_rect = canvas_cache.get("crop_rect") + if crop_rect is None and cov is not None: + crop_rect = compute_final_crop_rect(cov) + canvas_cache["crop_rect"] = crop_rect + if crop_rect is not None: + y0, y1, x0, x1 = crop_rect + roi = canvas[y0:y1, x0:x1] + else: + roi = canvas + + rh, rw = roi.shape[:2] + stitch_buf = canvas_cache.get("stitch_buf") + if stitch_buf is None or stitch_buf.shape[0] != rh or stitch_buf.shape[1] != rw: + stitch_buf = np.empty((rh, rw, 3), dtype=np.uint8) + canvas_cache["stitch_buf"] = stitch_buf + np.copyto(stitch_buf, roi) + t_after_finalize = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + if do_profile: + xmf = canvas_cache.get("x_min_f", 0.0) + ymf = canvas_cache.get("y_min_f", 0.0) + print( + f"stitch_feature_optimized: canvas out_w={out_w}, out_h={out_h}, " + f"x_min={xmf:.2f}, y_min={ymf:.2f} (fast_paste)" + ) + print( + "stitch_feature_optimized steps: " + f"brightness={t_after_brightness - t_profile0:.6f}s, " + f"geom={t_after_geom - t_after_brightness:.6f}s ({'hit' if geom_cache_hit else 'miss'}), " + f"reset={t_after_reset - t_after_geom:.6f}s, " + f"paste={t_after_paste - t_after_reset:.6f}s ({pasted_tiles + 1} tiles incl ref), " + f"finalize={t_after_finalize - t_after_paste:.6f}s, " + f"total={t_after_finalize - t_profile0:.6f}s" + ) + canvas_cache["last_coverage"] = coverage_ratio(cov) if cov is not None else canvas_cache.get("last_coverage", 0.0) + return stitch_buf + + +def render_feature_tiles( + t_ref_bgr, + tiles_bgr, + H_use_list, + *, + blend_backend: BlendBackend, + pad=5, + frame_idx=0, + print_every=30, + verbose=False, + match_brightness=False, + canvas_cache=None, + blend=True, +): + if H_use_list is None: + return None + + if canvas_cache is None: + canvas_cache = {} + do_profile = verbose and frame_idx % print_every == 0 + t_profile0 = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + if t_ref_bgr.ndim == 2: + t_ref_bgr = cv2.cvtColor(t_ref_bgr, cv2.COLOR_GRAY2BGR) + h0c, w0c = t_ref_bgr.shape[:2] + if match_brightness: + ref_mean_gray = mean_gray_bgr(t_ref_bgr) + tiles_bgr_c = [ + match_bgr_brightness( + tb if tb.ndim == 3 else cv2.cvtColor(tb, cv2.COLOR_GRAY2BGR), + ref_mean_gray, + ) + for tb in tiles_bgr + ] + elif any(tb.ndim == 2 for tb in tiles_bgr): + tiles_bgr_c = [tb if tb.ndim == 3 else cv2.cvtColor(tb, cv2.COLOR_GRAY2BGR) for tb in tiles_bgr] + else: + tiles_bgr_c = tiles_bgr + t_after_brightness = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + geom_cache_hit = ensure_feature_geom_cache( + canvas_cache, + t_ref_bgr, + tiles_bgr_c, + H_use_list, + pad, + allocate_blend_bufs=bool(blend), + ) + t_after_geom = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + if not blend: + return paste_feature_tiles_fast( + t_ref_bgr, + tiles_bgr_c, + H_use_list, + canvas_cache, + geom_cache_hit=geom_cache_hit, + frame_idx=frame_idx, + print_every=print_every, + verbose=verbose, + t_profile0=t_profile0, + t_after_brightness=t_after_brightness, + t_after_geom=t_after_geom, + ) + + out_w = canvas_cache["out_w"] + out_h = canvas_cache["out_h"] + ref_y0 = canvas_cache["ref_y0"] + ref_x0 = canvas_cache["ref_x0"] + pastes = canvas_cache["pastes"] + + blend_buf = canvas_cache["blend_buf"] + wsum = canvas_cache["wsum_buf"] + cov = canvas_cache["cov_buf"] + weight_maps = canvas_cache["weight_maps"] + if geom_cache_hit: + _dr = canvas_cache.get("dirty_rect") + if _dr is not None: + _dy0, _dy1, _dx0, _dx1 = _dr + blend_buf[_dy0:_dy1, _dx0:_dx1] = 0 + wsum[_dy0:_dy1, _dx0:_dx1] = 0 + cov[_dy0:_dy1, _dx0:_dx1] = 0 + else: + blend_buf.fill(0) + wsum.fill(0) + cov.fill(0) + t_after_reset = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + blend_margin = 16 + blended_tiles = 0 + for tile_bgr, H_use, paste in zip(tiles_bgr_c, H_use_list, pastes): + if H_use is None or paste is None: + continue + y0, x0 = paste + th, tw = tile_bgr.shape[:2] + wkey = (th, tw, blend_margin) + if wkey not in weight_maps: + weight_maps[wkey] = tile_weight_map(th, tw, blend_margin) + blend_backend.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] = tile_weight_map(h0c, w0c, blend_margin) + blend_backend.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") + if crop_rect is None: + crop_rect = compute_final_crop_rect(cov) + canvas_cache["crop_rect"] = crop_rect + if crop_rect is not None: + y0, y1, x0, x1 = crop_rect + roi = blend_buf[y0:y1, x0:x1] + else: + roi = blend_buf + rh, rw = roi.shape[:2] + stitch_buf = canvas_cache.get("stitch_buf") + if stitch_buf is None or stitch_buf.shape[0] != rh or stitch_buf.shape[1] != rw: + stitch_buf = np.empty((rh, rw, 3), dtype=np.uint8) + canvas_cache["stitch_buf"] = stitch_buf + cv2.convertScaleAbs(roi, dst=stitch_buf) + t_after_finalize = cv2.getTickCount() / cv2.getTickFrequency() if do_profile else 0.0 + + if do_profile: + xmf = canvas_cache.get("x_min_f", 0.0) + ymf = canvas_cache.get("y_min_f", 0.0) + print( + f"stitch_feature_optimized: canvas out_w={out_w}, out_h={out_h}, " + f"x_min={xmf:.2f}, y_min={ymf:.2f}" + ) + print( + "stitch_feature_optimized steps: " + f"brightness={t_after_brightness - t_profile0:.6f}s, " + f"geom={t_after_geom - t_after_brightness:.6f}s ({'hit' if geom_cache_hit else 'miss'}), " + f"reset={t_after_reset - t_after_geom:.6f}s, " + f"blend={t_after_blend - t_after_reset:.6f}s ({blended_tiles + 1} tiles incl ref), " + f"finalize={t_after_finalize - t_after_blend:.6f}s, " + f"total={t_after_finalize - t_profile0:.6f}s" + ) + + canvas_cache["last_coverage"] = coverage_ratio(cov) + canvas_cache["last_ms"] = ( + (t_after_finalize - t_profile0) * 1000.0 if do_profile else canvas_cache.get("last_ms", 0.0) + ) + return stitch_buf + + +def stitch_feature_optimized( + *, + isstitch: bool, + index, + nrow: int, + ncol: int, + tile_list: Sequence, + cache: Dict[str, Any], + blend_backend: BlendBackend, + blend: bool = True, + match_brightness: bool = False, + pad: int = 5, + print_every: int = 30, + verbose: bool = False, + frame_idx_holder: Optional[Dict[str, int]] = None, +): + """Estimate transforms once and composite the current tile grid.""" + plan = get_feature_stitch_plan(isstitch, index, nrow, ncol, tile_list) + if plan is None: + return None + + ref_idx, t_ref, moving_indices, moving_tiles, selected_positions = plan + if frame_idx_holder is None: + frame_idx_holder = {"idx": 0} + frame_idx_holder["idx"] = int(frame_idx_holder.get("idx", 0)) + 1 + frame_idx = frame_idx_holder["idx"] + tile_shapes = tuple(tile_list[pos].shape[:2] for pos in selected_positions) + cache_changed = ( + cache["ref_idx"] != ref_idx + or cache["moving_indices"] != moving_indices + or cache["tile_shapes"] != tile_shapes + ) + + if cache_changed: + cache["ref_idx"] = ref_idx + cache["moving_indices"] = moving_indices + cache["tile_shapes"] = tile_shapes + cache["transforms"] = None + cache["canvas"] = {} + if verbose: + print( + "Feature stitch tile selection: " + f"ref_idx={ref_idx}, moving_indices={list(moving_indices)}" + ) + + H_use_list = cache["transforms"] + if H_use_list is None: + do_verbose_first = verbose and frame_idx % print_every == 0 + H_use_list = [] + for j, tile in enumerate(moving_tiles): + H_use_j, _ = estimate_bfmatcher_stitch_transform( + t_ref, + tile, + verbose=(do_verbose_first and j == 0), + ) + H_use_list.append(H_use_j) + cache["transforms"] = H_use_list + + if not any(H is not None for H in H_use_list): + return None + + return render_feature_tiles( + t_ref, + moving_tiles, + H_use_list, + blend_backend=blend_backend, + pad=pad, + frame_idx=frame_idx, + print_every=print_every, + verbose=verbose, + match_brightness=match_brightness, + canvas_cache=cache["canvas"], + blend=blend, + ) diff --git a/src/orisys/stitch/ops.py b/src/orisys/stitch/ops.py new file mode 100644 index 0000000..bc23a1f --- /dev/null +++ b/src/orisys/stitch/ops.py @@ -0,0 +1,136 @@ +"""Pure array compute helpers for stitching (backend-agnostic). + +These functions take/return NumPy arrays and must not depend on ``Stitcher`` +instance state so future Torch/CuPy backends can reuse the same call shapes. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import cv2 +import numpy as np + + +def tile_weight_map(h: int, w: int, margin: int = 16) -> np.ndarray: + """Linear ramp weight: 0 at edges, 1.0 at ``margin`` pixels inward.""" + wy = np.ones(h, dtype=np.float32) + wx = np.ones(w, dtype=np.float32) + m = min(margin, h // 2, w // 2) + if m > 0: + ramp = np.linspace(0.0, 1.0, m + 1, dtype=np.float32)[1:] + wy[:m] = ramp + wy[-m:] = ramp[::-1] + wx[:m] = ramp + wx[-m:] = ramp[::-1] + return wy[:, None] * wx[None, :] + + +def paste_tile_clipped(canvas, tile, y0: int, x0: int, cov=None) -> None: + """Paste tile into canvas with clipping; optionally mark coverage.""" + th, tw = tile.shape[:2] + H, W = canvas.shape[:2] + y0d, x0d = max(0, y0), max(0, x0) + y1d, x1d = min(H, y0 + th), min(W, x0 + tw) + if y0d >= y1d or x0d >= x1d: + return + sy0, sx0 = y0d - y0, x0d - x0 + sh, sw = y1d - y0d, x1d - x0d + np.copyto(canvas[y0d:y1d, x0d:x1d], tile[sy0 : sy0 + sh, sx0 : sx0 + sw]) + if cov is not None: + cov[y0d:y1d, x0d:x1d] = 1 + + +def copy_into_tile_buffer(src_view, dst, flip_x: bool = True, flip_y: bool = True) -> None: + """Flip/copy ``src_view`` into preallocated ``dst``.""" + if flip_x and flip_y: + cv2.flip(src_view, -1, dst) + elif flip_x: + cv2.flip(src_view, 1, dst) + elif flip_y: + cv2.flip(src_view, 0, dst) + else: + np.copyto(dst, src_view) + + +def build_blend_coef(ov_pos: int, dim: int, ndim: int) -> np.ndarray: + """1D/2D blend ramp for legacy overlap stitching.""" + blend_coef = np.linspace(1.0, 0.0, ov_pos, dtype=np.float32) + if dim == 1: + return blend_coef.reshape(1, ov_pos, 1) if ndim == 3 else blend_coef.reshape(1, ov_pos) + return blend_coef.reshape(ov_pos, 1, 1) if ndim == 3 else blend_coef.reshape(ov_pos, 1) + + +def mean_gray_bgr(bgr) -> float: + ch = cv2.mean(bgr) + return (ch[0] + ch[1] + ch[2]) / 3.0 + + +def match_bgr_brightness(tile_bgr, ref_mean_gray: float): + delta = int(round(ref_mean_gray - mean_gray_bgr(tile_bgr))) + if delta == 0: + return tile_bgr + lut = np.clip(np.arange(256, dtype=np.int16) + delta, 0, 255).astype(np.uint8) + return cv2.LUT(tile_bgr, lut) + + +def largest_fully_covered_rect(cov) -> Optional[Tuple[int, int, int, int]]: + if cov.size == 0: + return None + + mask = cov > 0 + if not np.any(mask): + return None + + h, w = mask.shape + heights = np.zeros(w, dtype=np.int32) + best_area = 0 + best_rect = None + for y in range(h): + heights = np.where(mask[y], heights + 1, 0) + stack = [] + for x in range(w + 1): + cur_h = heights[x] if x < w else 0 + start = x + while stack and stack[-1][1] > cur_h: + start0, h0 = stack.pop() + area = h0 * (x - start0) + if area > best_area: + best_area = area + best_rect = (y - h0 + 1, y + 1, start0, x) + start = start0 + if not stack or stack[-1][1] < cur_h: + stack.append((start, cur_h)) + return best_rect + + +def compute_final_crop_rect(cov) -> Optional[Tuple[int, int, int, int]]: + valid = cov > 0 + if not np.any(valid): + return None + + ys, xs = np.where(valid) + by0, by1 = int(ys.min()), int(ys.max()) + 1 + bx0, bx1 = int(xs.min()), int(xs.max()) + 1 + rect2 = largest_fully_covered_rect(cov[by0:by1, bx0:bx1]) + if rect2 is None: + return by0, by1, bx0, bx1 + + y0, y1, x0, x1 = rect2 + return by0 + y0, by0 + y1, bx0 + x0, bx0 + x1 + + +def coverage_ratio(cov) -> float: + if cov is None or cov.size == 0: + return 0.0 + return float(np.mean(cov > 0)) + + +def translation_from_H_to_ref(H) -> Tuple[float, float]: + if H is None: + return 0.0, 0.0 + H = np.asarray(H, dtype=np.float64) + if abs(H[2, 2]) > 1e-12: + H = H / H[2, 2] + pt = H @ np.array([[0.0], [0.0], [1.0]], dtype=np.float64) + return float(pt[0, 0]), float(pt[1, 0]) diff --git a/src/orisys/stitch/tiles.py b/src/orisys/stitch/tiles.py new file mode 100644 index 0000000..60c32be --- /dev/null +++ b/src/orisys/stitch/tiles.py @@ -0,0 +1,86 @@ +"""Tile extraction and indexing for compound-eye frames.""" + +from __future__ import annotations + +from typing import Iterable, List, Optional, Sequence, Tuple + +import numpy as np + +from .ops import copy_into_tile_buffer + + +def extract_units_from_boxes( + img, + boxes: Sequence[Tuple[int, int, int, int]], + *, + selected_index: Iterable[int], + units: Optional[list] = None, + flip_x: bool = True, + flip_y: bool = True, +) -> list: + """Crop/flip tiles from ``boxes`` into ``units`` buffers (reused when possible).""" + h, w = img.shape[:2] + selected = set(selected_index) + flipped_units = [] if units is None else units + + for i, (x0, y0, x1, y1) in enumerate(boxes): + if i not in selected: + continue + + x0 = max(0, int(x0)) + y0 = max(0, int(y0)) + x1 = min(w, int(x1)) + y1 = min(h, int(y1)) + if x1 <= x0 or y1 <= y0: + continue + + src_view = img[y0:y1, x0:x1] + if units is None: + tile = np.empty_like(src_view) + copy_into_tile_buffer(src_view, tile, flip_x=flip_x, flip_y=flip_y) + flipped_units.append((x0, y0, tile)) + else: + _, _, tile = flipped_units[i] + if tile.shape != src_view.shape or tile.dtype != src_view.dtype: + tile = np.empty_like(src_view) + copy_into_tile_buffer(src_view, tile, flip_x=flip_x, flip_y=flip_y) + flipped_units[i] = (x0, y0, tile) + + return flipped_units + + +def get_tile_by_index(units: Sequence, idx) -> List: + if isinstance(idx, int): + idx = [idx] + if len(idx) > len(units): + raise ValueError("Index exceeds number of tiles.") + tile_list = [] + for i in range(len(idx)): + _xi, _yi, tile = units[idx[i]] + tile_list.append(tile) + return tile_list + + +def get_tile_index_by_position(units: Sequence, x, y, tol: int = 20) -> Optional[int]: + i = 0 + for xi, yi, _tile in units: + if abs(int(x) - int(xi)) <= tol and abs(int(y) - int(yi)) <= tol: + return i + i += 1 + return None + + +def get_unit_pos(units: Sequence, idx) -> List[Tuple[int, int]]: + if isinstance(idx, int): + idx = [idx] + if len(idx) > len(units): + raise ValueError("Index exceeds number of vision units.") + return [(units[i][0], units[i][1]) for i in idx] + + +def normalize_index(index) -> Tuple[int, ...]: + if isinstance(index, int): + return (index,) + if isinstance(index, np.ndarray): + return tuple(int(i) for i in index.tolist()) + return tuple(int(i) for i in index) diff --git a/tests/test_stitch.py b/tests/test_stitch.py new file mode 100644 index 0000000..9c24762 --- /dev/null +++ b/tests/test_stitch.py @@ -0,0 +1,168 @@ +"""Unit tests for orisys.stitch (synthetic tiles, no camera).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from orisys.Stitcher import Stitcher +from orisys.stitch import ( + StitchStatus, + get_blend_backend, + paste_tile_clipped, + resolve_blend_mode, + tile_weight_map, +) +from orisys.stitch.backends.numpy_backend import blend_kernel_numpy +from orisys.stitch.config import numba_blend_available, validate_box_grid +from orisys.stitch.ops import compute_final_crop_rect, coverage_ratio +from orisys.frame_processors import CompoundEyeFrameProcessor + + +def _solid_tile(val: int, h=32, w=32): + return np.full((h, w, 3), val, dtype=np.uint8) + + +def test_tile_weight_map_edges_are_low(): + wmap = tile_weight_map(32, 32, margin=4) + assert wmap.shape == (32, 32) + assert wmap.dtype == np.float32 + assert float(wmap[0, 0]) < float(wmap[16, 16]) + assert float(wmap[16, 16]) == pytest.approx(1.0) + + +def test_paste_tile_clipped_and_coverage(): + canvas = np.zeros((40, 40, 3), dtype=np.uint8) + cov = np.zeros((40, 40), dtype=np.uint8) + tile = _solid_tile(200, 10, 10) + paste_tile_clipped(canvas, tile, 5, 5, cov=cov) + assert int(canvas[5, 5, 0]) == 200 + assert int(cov[5, 5]) == 1 + assert int(cov[0, 0]) == 0 + + +def test_blend_kernel_numpy_first_paint(): + blend = np.zeros((4, 4, 3), dtype=np.float32) + wsum = np.zeros((4, 4), dtype=np.float32) + cov = np.zeros((4, 4), dtype=np.uint8) + src = np.ones((4, 4, 3), dtype=np.float32) * 50 + weight = np.ones((4, 4), dtype=np.float32) + blend_kernel_numpy(blend, wsum, cov, src, weight) + assert np.allclose(blend, 50) + assert np.all(cov == 1) + + +def test_numpy_backend_blend_tile(): + backend = get_blend_backend("numpy") + blend = np.zeros((20, 20, 3), dtype=np.float32) + wsum = np.zeros((20, 20), dtype=np.float32) + cov = np.zeros((20, 20), dtype=np.uint8) + src = np.full((8, 8, 3), 40, dtype=np.uint8) + weight = tile_weight_map(8, 8, margin=2) + backend.blend_tile_inplace(blend, wsum, cov, src, weight, 2, 3) + assert float(blend[4, 5, 0]) > 0 + assert coverage_ratio(cov) > 0 + + +@pytest.mark.skipif(not numba_blend_available(), reason="numba blend unavailable") +def test_numba_numpy_backend_parity(): + numpy_b = get_blend_backend("numpy") + numba_b = get_blend_backend("numba") + shape = (16, 16) + src = np.random.randint(0, 255, (*shape, 3), dtype=np.uint8) + weight = tile_weight_map(*shape, margin=3) + + def run(backend): + blend = np.zeros((*shape, 3), dtype=np.float32) + wsum = np.zeros(shape, dtype=np.float32) + cov = np.zeros(shape, dtype=np.uint8) + backend.blend_tile_inplace(blend, wsum, cov, src, weight, 0, 0) + return blend, cov + + b0, c0 = run(numpy_b) + b1, c1 = run(numba_b) + assert np.allclose(b0, b1, atol=1e-2) + assert np.array_equal(c0, c1) + + +def test_resolve_blend_mode_default(monkeypatch): + monkeypatch.delenv("ORISYS_STITCH_BLEND", raising=False) + monkeypatch.delenv("ORISYS_STITCH_BLEND_BACKEND", raising=False) + mode, fell = resolve_blend_mode({}) + assert mode == "none" + assert fell is False + + +def test_validate_box_grid(): + assert validate_box_grid([(0, 0, 1, 1)] * 9, 3, 3) + assert not validate_box_grid([(0, 0, 1, 1)] * 8, 3, 3) + + +def test_compute_final_crop_rect(): + cov = np.zeros((20, 20), dtype=np.uint8) + cov[4:16, 4:16] = 1 + rect = compute_final_crop_rect(cov) + assert rect is not None + y0, y1, x0, x1 = rect + assert y1 > y0 and x1 > x0 + + +def _mini_config(): + return { + "min_area_frac": 0.01, + "max_area_frac": 0.05, + "max_aspect_ratio": 50, + "index": [0, 1, 2, 3, 4, 5, 6, 7, 8], + "margin": 10, + "nominal_ov_pos": 100, + "n": [3, 3], + "row_tol": 0.6, + "box_size": [32, 32], + "blend": "numpy", + "n_boxes_target": 9, + } + + +def test_stitcher_status_and_rt_with_synthetic_tiles(): + stitcher = Stitcher(isstitch=True, config_stitcher=_mini_config(), verbose=False) + # Preload tiles as if extract_tiles already ran. + tiles = [_solid_tile(30 + i * 10) for i in range(9)] + stitcher.tile_list = tiles + stitcher.units = [(0, 0, t) for t in tiles] + stitcher.boxes = [(i * 40, 0, i * 40 + 32, 32) for i in range(9)] + + status = stitcher.status + assert isinstance(status, StitchStatus) + assert status.blend_mode == "numpy" + assert status.expected_boxes == 9 + + # Feature matching on flat solid tiles usually fails → None is OK. + out = stitcher.stitch_rt_optimized() + assert out is None or (isinstance(out, np.ndarray) and out.ndim == 3) + + +def test_compound_eye_processor_calls_stitcher(): + class FakeStitcher: + def __init__(self): + self.calls = [] + + def get_stitch_unit(self, frame): + self.calls.append("get") + + def extract_tiles(self): + self.calls.append("extract") + + def stitch_rt_optimized(self): + self.calls.append("stitch") + return np.zeros((8, 8, 3), dtype=np.uint8) + + fake = FakeStitcher() + proc = CompoundEyeFrameProcessor(fake) + out = proc.process(np.zeros((64, 64, 3), dtype=np.uint8)) + assert fake.calls == ["get", "extract", "stitch"] + assert out.shape == (8, 8, 3) + + +def test_unknown_backend_raises(): + with pytest.raises(ValueError): + get_blend_backend("torch")