Coverage for src/dictk/warp.py: 100%
95 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-30 00:04 +0000
1"""Subpixel point tracking by warping an affine kernel: inverse
2compositional Gauss-Newton (IC-GN)."""
4import numpy as np
5from scipy.ndimage import map_coordinates, spline_filter
7from dictk import translation
8from dictk.correlation import WindowingMethod
9from dictk.image import PixelCoordinate, SubpixelCoordinate
12def _spline_coefficients(image: np.ndarray) -> np.ndarray:
13 """Cubic B-spline coefficients of `image`, computed once so every
14 per-point refinement can sample it with `prefilter=False`."""
15 return spline_filter(image.astype(np.float64), order=3)
18def _gradients(image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
19 """`image`'s `(d/dy, d/dx)` central-difference gradients."""
20 gradient_y, gradient_x = np.gradient(image.astype(np.float64))
21 return gradient_y, gradient_x
24def _translation_warp(
25 *, reference_point: PixelCoordinate, start: PixelCoordinate
26) -> np.ndarray:
27 """The rigid seed: a warp that moves `reference_point` to the
28 whole-pixel `start`, with all four gradient terms zero."""
29 return np.array(
30 [
31 [1.0, 0.0, start.x - reference_point.x],
32 [0.0, 1.0, start.y - reference_point.y],
33 [0.0, 0.0, 1.0],
34 ]
35 )
38def _shape_seeds(
39 *,
40 reference_coefficients: np.ndarray,
41 current_image: np.ndarray,
42 current_coefficients: np.ndarray,
43 reference_point: PixelCoordinate,
44 kernel_margin_width: int,
45 kernel_margin_height: int,
46 search_margin_width: int,
47 search_margin_height: int,
48 gradient_limit: float,
49 gradient_step: float,
50 seed_count: int,
51) -> list[np.ndarray]:
52 """Deformation-aware seeds for one point, best first.
54 Tries every displacement gradient `(u_x, u_y, v_x, v_y)` on a grid
55 from `-gradient_limit` to `gradient_limit` in steps of
56 `gradient_step`. Each trial deforms the reference kernel by its
57 gradient, then finds that kernel's whole-pixel translation by phase
58 correlation against the search area, as `translation.locate` does
59 for a rigid kernel. Each trial is scored by ZNSSD at that
60 translation. Returns the `seed_count` lowest-scoring trials as 3x3
61 warps, ready for `_refine`."""
62 x0, y0 = reference_point.x, reference_point.y
63 dy, dx = np.mgrid[
64 -kernel_margin_height : kernel_margin_height + 1,
65 -kernel_margin_width : kernel_margin_width + 1,
66 ].astype(np.float64)
67 count = int(round(gradient_limit / gradient_step))
68 values = np.arange(-count, count + 1) * gradient_step
69 gradients = np.stack(
70 np.meshgrid(values, values, values, values, indexing="ij"), axis=-1
71 ).reshape(-1, 4)
72 shapes = np.eye(2) + gradients.reshape(-1, 2, 2)
74 # The deformed kernel, on the current image's pixel grid: the
75 # reference sampled at each offset carried back through the shape.
76 inverses = np.linalg.inv(shapes)
77 xs = x0 + inverses[:, 0, 0, None, None] * dx + inverses[:, 0, 1, None, None] * dy
78 ys = y0 + inverses[:, 1, 0, None, None] * dx + inverses[:, 1, 1, None, None] * dy
79 kernels = map_coordinates(
80 reference_coefficients, [ys, xs], order=3, prefilter=False
81 )
83 # Phase correlation, with each kernel centered in a zero-padded copy
84 # the size of the search area, as translation.locate pads its own.
85 search = current_image[
86 y0 - search_margin_height : y0 + search_margin_height,
87 x0 - search_margin_width : x0 + search_margin_width,
88 ].astype(np.float64)
89 rows, columns = search.shape
90 top = (rows - kernels.shape[1]) // 2
91 left = (columns - kernels.shape[2]) // 2
92 padded = np.zeros((len(kernels), rows, columns))
93 padded[:, top : top + kernels.shape[1], left : left + kernels.shape[2]] = kernels
94 product = np.fft.fft2(search)[None] * np.fft.fft2(padded).conj()
95 product /= np.maximum(np.abs(product), 100 * np.finfo(np.float64).eps)
96 surfaces = np.abs(np.fft.ifft2(product)).reshape(len(kernels), -1)
97 peak_rows, peak_columns = np.unravel_index(surfaces.argmax(axis=1), (rows, columns))
98 peak_rows = np.where(peak_rows >= rows // 2, peak_rows - rows, peak_rows)
99 peak_columns = np.where(
100 peak_columns >= columns // 2, peak_columns - columns, peak_columns
101 )
102 # A zero shift puts the kernel's center at search index
103 # (top + margin, left + margin), which is the point's own position
104 # minus this offset.
105 u = peak_columns + left + kernel_margin_width - search_margin_width
106 v = peak_rows + top + kernel_margin_height - search_margin_height
108 xs = x0 + u[:, None, None] + shapes[:, 0, 0, None, None] * dx
109 xs += shapes[:, 0, 1, None, None] * dy
110 ys = y0 + v[:, None, None] + shapes[:, 1, 0, None, None] * dx
111 ys += shapes[:, 1, 1, None, None] * dy
112 g = map_coordinates(current_coefficients, [ys, xs], order=3, prefilter=False)
113 g = g - g.mean(axis=(1, 2), keepdims=True)
114 g /= np.maximum(np.sqrt((g**2).sum(axis=(1, 2), keepdims=True)), 1e-12)
115 f = map_coordinates(
116 reference_coefficients, [y0 + dy, x0 + dx], order=3, prefilter=False
117 )
118 f = f - f.mean()
119 f /= np.sqrt((f**2).sum())
120 znssds = ((g - f) ** 2).sum(axis=(1, 2))
122 seeds = []
123 for index in np.argsort(znssds)[:seed_count]:
124 seed = np.eye(3)
125 seed[:2, :2] = shapes[index]
126 seed[:2, 2] = u[index], v[index]
127 seeds.append(seed)
128 return seeds
131def _refine(
132 *,
133 reference_image: np.ndarray,
134 reference_gradients: tuple[np.ndarray, np.ndarray],
135 current_coefficients: np.ndarray,
136 reference_point: PixelCoordinate,
137 initial_warp: np.ndarray,
138 kernel_margin_width: int,
139 kernel_margin_height: int,
140 max_iterations: int,
141 tolerance: float,
142) -> tuple[np.ndarray, float, int]:
143 """IC-GN refinement of one point, starting from the 3x3
144 `initial_warp` -- see [`fit`](#fit) for the method and the returned
145 matrix. `_translation_warp` builds the rigid seed `fit` uses. Takes
146 precomputed gradients and spline coefficients so a batch caller
147 computes them once, not once per point.
149 Also returns the ZNSSD between the reference kernel and
150 `current_image` sampled through the returned warp, from 0 for a
151 perfect match to 4, and how many updates IC-GN applied. That count
152 equals `max_iterations` when IC-GN stopped without meeting
153 `tolerance`."""
154 gradient_y, gradient_x = reference_gradients
155 x0, y0 = reference_point.x, reference_point.y
156 rows = slice(y0 - kernel_margin_height, y0 + kernel_margin_height + 1)
157 columns = slice(x0 - kernel_margin_width, x0 + kernel_margin_width + 1)
158 dy, dx = np.mgrid[
159 -kernel_margin_height : kernel_margin_height + 1,
160 -kernel_margin_width : kernel_margin_width + 1,
161 ].astype(np.float64)
163 f = reference_image[rows, columns].astype(np.float64)
164 f = f - f.mean()
165 f_norm = np.sqrt((f**2).sum())
167 gx = gradient_x[rows, columns]
168 gy = gradient_y[rows, columns]
169 steepest_descent = np.stack(
170 [gx, gx * dx, gx * dy, gy, gy * dx, gy * dy], axis=-1
171 ).reshape(-1, 6)
172 hessian_inverse = np.linalg.inv(steepest_descent.T @ steepest_descent)
174 def residual_at(warp):
175 xs = x0 + warp[0, 0] * dx + warp[0, 1] * dy + warp[0, 2]
176 ys = y0 + warp[1, 0] * dx + warp[1, 1] * dy + warp[1, 2]
177 g = map_coordinates(current_coefficients, [ys, xs], order=3, prefilter=False)
178 g = g - g.mean()
179 return f_norm / np.sqrt((g**2).sum()) * g - f
181 warp = np.array(initial_warp, dtype=np.float64)
182 for iterations in range(1, max_iterations + 1):
183 residual = residual_at(warp)
184 dp = hessian_inverse @ (steepest_descent.T @ residual.ravel())
185 increment = np.array(
186 [
187 [1.0 + dp[1], dp[2], dp[0]],
188 [dp[4], 1.0 + dp[5], dp[3]],
189 [0.0, 0.0, 1.0],
190 ]
191 )
192 warp = warp @ np.linalg.inv(increment)
193 if np.abs(dp).max() < tolerance:
194 break
196 znssd = float((residual_at(warp) ** 2).sum() / f_norm**2)
197 return warp, znssd, iterations
200def fit(
201 *,
202 reference_image: np.ndarray,
203 current_image: np.ndarray,
204 reference_point: PixelCoordinate,
205 search_center: PixelCoordinate,
206 kernel_margin_width: int,
207 kernel_margin_height: int,
208 search_margin_width: int,
209 search_margin_height: int,
210 windowing: WindowingMethod | None = None,
211 max_iterations: int = 50,
212 tolerance: float = 1e-6,
213) -> np.ndarray:
214 """Fit the affine warp that carries a kernel of `reference_image`
215 onto `current_image`.
217 Two stages. First,
218 [`dictk.translation.locate`](./translation.html#locate) finds the
219 nearest whole-pixel position. Second, inverse compositional
220 Gauss-Newton (IC-GN) refines it. IC-GN deforms a square kernel of
221 `reference_image`, centered on `reference_point`, by a six-parameter
222 affine warp. It then adjusts that warp until the warped kernel best
223 matches `current_image`, by zero-normalized sum of squared
224 differences (ZNSSD). Cubic B-spline interpolation samples
225 `current_image` at the warped, fractional positions.
227 The whole-pixel stage still matches a rigid, unwarped kernel. A
228 deformation large enough to defeat that match, such as a stretch of
229 about 10% or more with shear on a fine speckle pattern, leaves IC-GN
230 starting too far away to converge. `fit` has no check for this.
232 See [Kernel Warping](../getting_started/kernel_warping.html) for why
233 this beats phase correlation's upsampled peak, and
234 [`locate`](#locate) for the tracked position alone.
236 This kernel is `2 * kernel_margin_width + 1` wide and
237 `2 * kernel_margin_height + 1` tall, centered exactly on
238 `reference_point`. That is one pixel larger than `translation.locate`'s
239 own even kernel in each direction.
241 Args:
242 reference_image: The reference (undeformed) 2D grayscale image.
243 current_image: The current (deformed) 2D grayscale image.
244 reference_point: The point's fixed, known position, in
245 `reference_image`'s pixel reference frame.
246 search_center: Where to center the whole-pixel search, in
247 `current_image`'s pixel reference frame -- see
248 `translation.locate`'s own docstring.
249 kernel_margin_width: Half the kernel's width, in pixels, not
250 counting its center column. Must be >= 1.
251 kernel_margin_height: Half the kernel's height, in pixels, not
252 counting its center row. Must be >= 1.
253 search_margin_width: Half the whole-pixel search area's width,
254 in pixels. Must be greater than `kernel_margin_width`.
255 search_margin_height: Half the whole-pixel search area's height,
256 in pixels. Must be greater than `kernel_margin_height`.
257 windowing: Passed straight through to the whole-pixel
258 `translation.locate` stage. Default `None` applies no
259 windowing. IC-GN itself never windows.
260 max_iterations: The most IC-GN updates to apply. Default `50`.
261 Must be >= 1.
262 tolerance: IC-GN stops once every warp parameter's update falls
263 below this. Default `1e-6`. Must be > 0.
265 Returns:
266 The 3x3 warp matrix
268 [[1 + u_x, u_y, u],
269 [v_x, 1 + v_y, v],
270 [0, 0, 1]]
272 It maps a pixel's offset `(dx, dy, 1)` from `reference_point`
273 in `reference_image` to that pixel's offset from
274 `reference_point` in `current_image`. `(u, v)` is the tracked
275 displacement. `u_x`, `u_y`, `v_x`, `v_y` are the displacement
276 gradients across the kernel.
278 Raises:
279 ValueError: If any margin fails `translation.locate`'s own
280 checks, `max_iterations` is less than 1, or `tolerance` is
281 not positive.
282 """
283 if max_iterations < 1:
284 raise ValueError(f"max_iterations {max_iterations} must be >= 1")
285 if tolerance <= 0:
286 raise ValueError(f"tolerance {tolerance} must be > 0")
288 start = translation.locate(
289 reference_image=reference_image,
290 current_image=current_image,
291 reference_point=reference_point,
292 search_center=search_center,
293 kernel_margin_width=kernel_margin_width,
294 kernel_margin_height=kernel_margin_height,
295 search_margin_width=search_margin_width,
296 search_margin_height=search_margin_height,
297 windowing=windowing,
298 )
299 fitted, _znssd, _iterations = _refine(
300 reference_image=reference_image,
301 reference_gradients=_gradients(reference_image),
302 current_coefficients=_spline_coefficients(current_image),
303 reference_point=reference_point,
304 initial_warp=_translation_warp(reference_point=reference_point, start=start),
305 kernel_margin_width=kernel_margin_width,
306 kernel_margin_height=kernel_margin_height,
307 max_iterations=max_iterations,
308 tolerance=tolerance,
309 )
310 return fitted
313def locate(
314 *,
315 reference_image: np.ndarray,
316 current_image: np.ndarray,
317 reference_point: PixelCoordinate,
318 search_center: PixelCoordinate,
319 kernel_margin_width: int,
320 kernel_margin_height: int,
321 search_margin_width: int,
322 search_margin_height: int,
323 windowing: WindowingMethod | None = None,
324 max_iterations: int = 50,
325 tolerance: float = 1e-6,
326) -> SubpixelCoordinate:
327 """Subpixel point tracking by warping an affine kernel, free of the
328 pixel-locking bias of
329 [`dictk.translation.locate_subpixel`](./translation.html#locate_subpixel).
331 Runs [`fit`](#fit) and returns only the fitted warp's translation,
332 as a position. Every argument passes straight through to `fit` --
333 see its own docstring for each one, and for the method.
335 Returns:
336 The point's location, in `current_image`'s pixel reference
337 frame, generally fractional.
339 Raises:
340 ValueError: Under the same conditions as `fit`.
341 """
342 warp = fit(
343 reference_image=reference_image,
344 current_image=current_image,
345 reference_point=reference_point,
346 search_center=search_center,
347 kernel_margin_width=kernel_margin_width,
348 kernel_margin_height=kernel_margin_height,
349 search_margin_width=search_margin_width,
350 search_margin_height=search_margin_height,
351 windowing=windowing,
352 max_iterations=max_iterations,
353 tolerance=tolerance,
354 )
355 return SubpixelCoordinate(
356 x=reference_point.x + warp[0, 2], y=reference_point.y + warp[1, 2]
357 )