Spaces:
Sleeping
Sleeping
File size: 7,536 Bytes
aa4d14b 99e1f27 aa4d14b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | """
Shared utilities for model loading, tile processing, and multi-scale detection.
"""
import cv2
import numpy as np
def make_blend_weights(tile_size: int, overlap: int) -> np.ndarray:
"""Raised-cosine (trapezoid) 2D weight for seamless tile stitching.
Ramps from 0 to 1 across the overlap band on each edge and stays flat in
the interior, so overlapping tile predictions blend without visible seams.
"""
overlap = max(0, min(overlap, tile_size // 2))
if overlap == 0:
return np.ones((tile_size, tile_size), dtype=np.float32)
ramp = np.linspace(0.0, 1.0, overlap, dtype=np.float32)
flat = np.ones(tile_size - 2 * overlap, dtype=np.float32)
profile = np.concatenate([ramp, flat, ramp[::-1]])
return np.outer(profile, profile).astype(np.float32)
def tiled_score_map(score_tile_fn, img1, img2, tile_size: int = 256,
overlap: int | None = None,
score_batch_fn=None, batch: int = 1):
"""Sliding-window tiled scoring with reflect padding + cosine blending.
``score_tile_fn(tile1, tile2)`` must return a float32 array in [0, 1] the
same height/width as the input tile. Returns a float32 change-probability
map cropped back to the original (h, w). This is the single shared
implementation used by every deep model so blending stays consistent.
When ``batch > 1`` and ``score_batch_fn`` is provided, tiles are scored in
groups (``score_batch_fn(list[(t1, t2)]) -> list[prob]``) so GPUs can run
several tiles per forward pass.
"""
h, w = img1.shape[:2]
if overlap is None:
overlap = tile_size // 4
overlap = max(0, min(overlap, tile_size // 2))
stride = max(1, tile_size - overlap)
pad_h = (tile_size - h % tile_size) % tile_size
pad_w = (tile_size - w % tile_size) % tile_size
# Guarantee at least one full tile even for small inputs
pad_h = max(pad_h, max(0, tile_size - h))
pad_w = max(pad_w, max(0, tile_size - w))
if pad_h or pad_w:
img1 = np.pad(img1, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
img2 = np.pad(img2, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
ph, pw = img1.shape[:2]
score_sum = np.zeros((ph, pw), dtype=np.float32)
count = np.zeros((ph, pw), dtype=np.float32)
weight_2d = make_blend_weights(tile_size, overlap)
ys = list(range(0, ph - tile_size + 1, stride))
xs = list(range(0, pw - tile_size + 1, stride))
# Ensure the far edges are always covered
if ys and ys[-1] != ph - tile_size:
ys.append(ph - tile_size)
if xs and xs[-1] != pw - tile_size:
xs.append(pw - tile_size)
coords = [(y0, x0) for y0 in ys for x0 in xs]
def _accumulate(y0, x0, prob):
if prob.shape != (tile_size, tile_size):
prob = cv2.resize(prob.astype(np.float32), (tile_size, tile_size),
interpolation=cv2.INTER_LINEAR)
score_sum[y0:y0 + tile_size, x0:x0 + tile_size] += prob * weight_2d
count[y0:y0 + tile_size, x0:x0 + tile_size] += weight_2d
use_batch = score_batch_fn is not None and batch > 1
if use_batch:
for i in range(0, len(coords), batch):
chunk = coords[i:i + batch]
pairs = [
(np.ascontiguousarray(img1[y:y + tile_size, x:x + tile_size]),
np.ascontiguousarray(img2[y:y + tile_size, x:x + tile_size]))
for (y, x) in chunk
]
probs = score_batch_fn(pairs)
for (y0, x0), prob in zip(chunk, probs):
_accumulate(y0, x0, prob)
else:
for (y0, x0) in coords:
t1 = np.ascontiguousarray(img1[y0:y0 + tile_size, x0:x0 + tile_size])
t2 = np.ascontiguousarray(img2[y0:y0 + tile_size, x0:x0 + tile_size])
_accumulate(y0, x0, score_tile_fn(t1, t2))
count = np.maximum(count, 1e-6)
avg = score_sum / count
return avg[:h, :w]
def split_into_tiles(img, tile_size=512, overlap=64):
"""
Split an image into overlapping tiles.
Returns list of (tile, y_offset, x_offset) tuples.
"""
h, w = img.shape[:2]
stride = tile_size - overlap
tiles = []
pad_h = (tile_size - h % tile_size) % tile_size if h % tile_size else 0
pad_w = (tile_size - w % tile_size) % tile_size if w % tile_size else 0
if pad_h or pad_w:
img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)) if img.ndim == 3
else ((0, pad_h), (0, pad_w)), mode="reflect")
ph, pw = img.shape[:2]
for y in range(0, ph - tile_size + 1, stride):
for x in range(0, pw - tile_size + 1, stride):
tile = img[y:y+tile_size, x:x+tile_size]
tiles.append((tile, y, x))
return tiles, (ph, pw), (h, w)
def merge_tile_masks(tile_results, padded_shape, orig_shape, tile_size=512, overlap=64):
"""
Merge tile-level binary masks back into a single full-resolution mask.
Uses raised-cosine blending to avoid tile boundary artifacts.
"""
ph, pw = padded_shape
h, w = orig_shape
score_sum = np.zeros((ph, pw), dtype=np.float32)
count = np.zeros((ph, pw), dtype=np.float32)
ramp = np.linspace(0, 1, overlap)
flat = np.ones(tile_size - 2 * overlap)
profile = np.concatenate([ramp, flat, ramp[::-1]])
weight_2d = np.outer(profile, profile).astype(np.float32)
for (mask_tile, y, x) in tile_results:
score = mask_tile.astype(np.float32) / 255.0 if mask_tile.max() > 1 else mask_tile.astype(np.float32)
if score.shape != (tile_size, tile_size):
score = cv2.resize(score, (tile_size, tile_size))
score_sum[y:y+tile_size, x:x+tile_size] += score * weight_2d
count[y:y+tile_size, x:x+tile_size] += weight_2d
count = np.maximum(count, 1e-6)
merged = score_sum / count
merged = merged[:h, :w]
return (merged * 255).astype(np.uint8)
def multiscale_detect(detect_fn, img1, img2, scales=(1.0, 0.5, 0.25)):
"""
Run a detection function at multiple scales and combine via logical OR.
detect_fn(img1, img2) -> uint8 mask [0|255].
Captures small structures at full res and large regions at coarse scales.
"""
h, w = img1.shape[:2]
combined = np.zeros((h, w), dtype=np.uint8)
for scale in scales:
if scale == 1.0:
s1, s2 = img1, img2
else:
sh, sw = max(64, int(h * scale)), max(64, int(w * scale))
s1 = cv2.resize(img1, (sw, sh), interpolation=cv2.INTER_AREA)
s2 = cv2.resize(img2, (sw, sh), interpolation=cv2.INTER_AREA)
mask = detect_fn(s1, s2)
if scale != 1.0:
mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
combined = np.maximum(combined, mask)
return combined
def build_confidence_map(channels, weights=None):
"""
Build a [0-1] confidence map from multiple normalized signal channels.
Each channel should be a float32 array in [0,1].
If weights is None, uses equal weighting.
"""
if not channels:
return None
if weights is None:
weights = [1.0 / len(channels)] * len(channels)
total_w = sum(weights)
weights = [w / total_w for w in weights]
shape = channels[0].shape
conf = np.zeros(shape, dtype=np.float64)
for ch, w in zip(channels, weights):
if ch.shape != shape:
ch = cv2.resize(ch.astype(np.float32), (shape[1], shape[0]))
conf += w * ch.astype(np.float64)
return np.clip(conf, 0, 1).astype(np.float32)
|