diff --git a/pm2d/_jit_kernels.py b/pm2d/_jit_kernels.py index 626a346..b504819 100644 --- a/pm2d/_jit_kernels.py +++ b/pm2d/_jit_kernels.py @@ -328,6 +328,65 @@ if HAS_NUMBA: out[vi] = best if best > 0.0 else 0.0 return out + @nb.njit(cache=True, parallel=True, fastmath=True, boundscheck=False) + def _jit_score_bitmap_rescored_u16( + spread: np.ndarray, # uint16 (H, W) - 16 bit di polarity-aware + dx: np.ndarray, dy: np.ndarray, bins: np.ndarray, + bit_active: np.uint16, + bg: np.ndarray, + ) -> np.ndarray: + """Versione uint16 di _jit_score_bitmap_rescored per polarity 16-bin. + + Identica logica ma mask = uint16(1) << b dove b in [0..15] + (orientamento mod 2π invece di mod π). + """ + H, W = spread.shape + N = dx.shape[0] + acc = np.zeros((H, W), dtype=np.float32) + for y in nb.prange(H): + for i in range(N): + b = bins[i] + mask = np.uint16(1) << b + if (bit_active & mask) == 0: + continue + ddy = dy[i] + yy = y + ddy + if yy < 0 or yy >= H: + continue + ddx = dx[i] + x_lo = 0 if ddx >= 0 else -ddx + x_hi = W if ddx <= 0 else W - ddx + for x in range(x_lo, x_hi): + if spread[yy, x + ddx] & mask: + acc[y, x] += 1.0 + if N > 0: + inv = 1.0 / N + for y in nb.prange(H): + for x in range(W): + v = acc[y, x] * inv + bgv = bg[y, x] + if bgv < 1.0: + r = (v - bgv) / (1.0 - bgv + 1e-6) + acc[y, x] = r if r > 0.0 else 0.0 + else: + acc[y, x] = 0.0 + return acc + + @nb.njit(cache=True, parallel=True, fastmath=True, boundscheck=False) + def _jit_popcount_density_u16(spread: np.ndarray) -> np.ndarray: + """Popcount per uint16 (16 bin polarity).""" + H, W = spread.shape + out = np.zeros((H, W), dtype=np.float32) + for y in nb.prange(H): + for x in range(W): + v = spread[y, x] + cnt = 0 + for b in range(16): + if v & (np.uint16(1) << b): + cnt += 1 + out[y, x] = float(cnt) + return out + @nb.njit(cache=True, parallel=True, fastmath=True, boundscheck=False) def _jit_popcount_density(spread: np.ndarray) -> np.ndarray: """Conta bit set per pixel: ritorna (H, W) float32 in [0..8].""" @@ -368,6 +427,11 @@ if HAS_NUMBA: spread, dx, dy, b, offsets, np.uint8(0xFF), bg_pv, scale_idx, ) _jit_popcount_density(spread) + spread16 = np.zeros((32, 32), dtype=np.uint16) + _jit_score_bitmap_rescored_u16( + spread16, dx, dy, b, np.uint16(0xFFFF), bg, + ) + _jit_popcount_density_u16(spread16) else: # pragma: no cover @@ -392,6 +456,12 @@ else: # pragma: no cover ): raise RuntimeError("numba non disponibile") + def _jit_score_bitmap_rescored_u16(spread, dx, dy, bins, bit_active, bg): + raise RuntimeError("numba non disponibile") + + def _jit_popcount_density_u16(spread): + raise RuntimeError("numba non disponibile") + def _jit_popcount_density(spread): raise RuntimeError("numba non disponibile") @@ -426,16 +496,20 @@ def score_bitmap_rescored( ) -> np.ndarray: """Score bitmap + rescore fusi in un solo pass (JIT). - stride > 1: valuta solo pixel su griglia stride×stride. Le celle non - valutate restano 0 nello score map. Pensato per coarse-pass al top - della piramide; il refinement full-res poi recupera precisione. + Dispatch per dtype: uint16 → kernel polarity 16-bin, uint8 → kernel + standard 8-bin (con eventuale stride > 1 per coarse top-level). """ if HAS_NUMBA and len(dx) > 0: - spread_c = np.ascontiguousarray(spread, dtype=np.uint8) dx_c = np.ascontiguousarray(dx, dtype=np.int32) dy_c = np.ascontiguousarray(dy, dtype=np.int32) bins_c = np.ascontiguousarray(bins, dtype=np.int8) bg_c = np.ascontiguousarray(bg, dtype=np.float32) + if spread.dtype == np.uint16: + spread_c = np.ascontiguousarray(spread, dtype=np.uint16) + return _jit_score_bitmap_rescored_u16( + spread_c, dx_c, dy_c, bins_c, np.uint16(bit_active), bg_c, + ) + spread_c = np.ascontiguousarray(spread, dtype=np.uint8) if stride > 1: return _jit_score_bitmap_rescored_strided( spread_c, dx_c, dy_c, bins_c, np.uint8(bit_active), bg_c, @@ -528,6 +602,17 @@ def popcount_density(spread: np.ndarray) -> np.ndarray: 2) numpy.bitwise_count (NumPy 2.0+, SIMD ma single-thread) 3) Fallback numpy bit-shift puro """ + if spread.dtype == np.uint16: + spread_c = np.ascontiguousarray(spread, dtype=np.uint16) + if HAS_NUMBA: + return _jit_popcount_density_u16(spread_c) + if _HAS_NP_BITCOUNT: + return np.bitwise_count(spread_c).astype(np.float32, copy=False) + H, W = spread_c.shape + out = np.zeros((H, W), dtype=np.float32) + for b in range(16): + out += ((spread_c >> b) & 1).astype(np.float32) + return out spread_c = np.ascontiguousarray(spread, dtype=np.uint8) if HAS_NUMBA: return _jit_popcount_density(spread_c) diff --git a/pm2d/line_matcher.py b/pm2d/line_matcher.py index ece8db4..d3a1f79 100644 --- a/pm2d/line_matcher.py +++ b/pm2d/line_matcher.py @@ -46,7 +46,8 @@ from pm2d._jit_kernels import ( HAS_NUMBA, ) -N_BINS = 8 # orientamenti quantizzati modulo π +N_BINS = 8 # default: orientamento mod π (no polarity) +N_BINS_POL = 16 # use_polarity=True: orientamento mod 2π (con polarity) def _poly_iou(p1: np.ndarray, p2: np.ndarray) -> float: @@ -143,6 +144,7 @@ class LineShapeMatcher: pyramid_levels: int = 2, top_score_factor: float = 0.5, n_threads: int | None = None, + use_polarity: bool = False, ) -> None: self.num_features = num_features self.weak_grad = weak_grad @@ -156,6 +158,12 @@ class LineShapeMatcher: self.pyramid_levels = max(1, pyramid_levels) self.top_score_factor = top_score_factor self.n_threads = n_threads or max(1, (os.cpu_count() or 2) - 1) + # Polarity-aware: 16 bin (orientamento mod 2π) usando bitmap uint16. + # Distingue edge "chiaro→scuro" da "scuro→chiaro" → 2x selettività. + # Usare quando background di scena varia (chiaro/scuro) e orientamento + # template e' direzionale. + self.use_polarity = use_polarity + self._n_bins = N_BINS_POL if use_polarity else N_BINS self.variants: list[_Variant] = [] self.template_size: tuple[int, int] = (0, 0) @@ -171,15 +179,20 @@ class LineShapeMatcher: return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) return img - @staticmethod - def _gradient(gray: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + def _gradient(self, gray: np.ndarray) -> tuple[np.ndarray, np.ndarray]: gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3) gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3) mag = cv2.magnitude(gx, gy) - ang = np.arctan2(gy, gx) - ang_mod = np.where(ang < 0, ang + np.pi, ang) - bins = np.floor(ang_mod / np.pi * N_BINS).astype(np.int16) - bins = np.clip(bins, 0, N_BINS - 1) + ang = np.arctan2(gy, gx) # [-π, π] + if self.use_polarity: + # Mod 2π: bin 0..15 codifica direzione + polarity edge. + ang_full = np.where(ang < 0, ang + 2.0 * np.pi, ang) + bins = np.floor(ang_full / (2.0 * np.pi) * N_BINS_POL).astype(np.int16) + bins = np.clip(bins, 0, N_BINS_POL - 1) + else: + ang_mod = np.where(ang < 0, ang + np.pi, ang) + bins = np.floor(ang_mod / np.pi * N_BINS).astype(np.int16) + bins = np.clip(bins, 0, N_BINS - 1) return mag, bins def _extract_features( @@ -390,20 +403,22 @@ class LineShapeMatcher: return raw def _spread_bitmap(self, gray: np.ndarray) -> np.ndarray: - """Spread bitmap uint8: bit b acceso dove bin b è presente nel raggio. + """Spread bitmap: bit b acceso dove bin b è presente nel raggio. - Formato compatto 32× più denso della response map (N_BINS, H, W) float32. + dtype: uint8 per N_BINS=8, uint16 per N_BINS_POL=16 (use_polarity). """ mag, bins = self._gradient(gray) valid = mag >= self.weak_grad k = 2 * self.spread_radius + 1 kernel = np.ones((k, k), dtype=np.uint8) H, W = gray.shape - spread = np.zeros((H, W), dtype=np.uint8) - for b in range(N_BINS): + nb = self._n_bins + dtype = np.uint16 if nb > 8 else np.uint8 + spread = np.zeros((H, W), dtype=dtype) + for b in range(nb): mask_b = ((bins == b) & valid).astype(np.uint8) d = cv2.dilate(mask_b, kernel) - spread |= (d << b) + spread |= (d.astype(dtype) << b) return spread @staticmethod @@ -653,9 +668,10 @@ class LineShapeMatcher: x_lo = int(cx) - margin; x_hi = int(cx) + margin + 1 sh_w = y_hi - y_lo; sw_w = x_hi - x_lo acc = np.zeros((sh_w, sw_w), dtype=np.float32) + spread_dtype = spread0.dtype.type for i in range(len(dx)): ddx = int(dx[i]); ddy = int(dy[i]); b = int(fb[i]) - bit = np.uint8(1 << b) + bit = spread_dtype(1 << b) sy0 = y_lo + ddy; sy1 = y_hi + ddy sx0 = x_lo + ddx; sx1 = x_hi + ddx a_y0 = max(0, -sy0); a_y1 = sh_w - max(0, sy1 - H) @@ -824,8 +840,8 @@ class LineShapeMatcher: # map float32 → MOLTO più cache-friendly per _score_by_shift. spread_top = self._spread_bitmap(grays[top]) bit_active_top = int( - sum(1 << b for b in range(N_BINS) - if (spread_top & np.uint8(1 << b)).any()) + sum(1 << b for b in range(self._n_bins) + if (spread_top & (spread_top.dtype.type(1) << b)).any()) ) if nms_radius is None: nms_radius = max(8, min(self.template_size) // 2) @@ -982,8 +998,8 @@ class LineShapeMatcher: # Full-res (parallelizzato) con bitmap spread0 = self._spread_bitmap(gray0) bit_active_full = int( - sum(1 << b for b in range(N_BINS) - if (spread0 & np.uint8(1 << b)).any()) + sum(1 << b for b in range(self._n_bins) + if (spread0 & (spread0.dtype.type(1) << b)).any()) ) density_full = _jit_popcount(spread0) for sc in unique_scales: