Coverage for src/dictk/grid.py: 98%
166 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-30 00:04 +0000
1"""A rectangular collection of points, and batch point tracking across it."""
3import heapq
4from collections.abc import Sequence
5from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
6from enum import Enum
7from functools import partial
8from typing import Final
10import numpy as np
11from scipy.spatial import cKDTree
13from dictk import translation, warp
14from dictk.correlation import WindowingMethod
15from dictk.image import PixelCoordinate, SubpixelCoordinate
17# How many nearest points `locate_warp_propagated` links each point to.
18# On a rectangular grid, 8 covers the 4 edge neighbors and 4 diagonals.
19NEIGHBOR_COUNT: Final[int] = 8
21# How many of the start search's best-scoring trials IC-GN refines.
22START_SEED_COUNT: Final[int] = 3
24# How many start searches in a row may fail `znssd_limit` before
25# `locate_warp_propagated` stops trying new start points. A 40x40 px
26# patch of unrelated speckle over the image center needed more than 5.
27START_ATTEMPT_LIMIT: Final[int] = 10
30class Executor(Enum):
31 """Pool type `locate`'s `max_workers` parameter runs on.
33 - THREAD: shares memory with the caller, no pickling. Pays a small
34 scheduling cost on every task, which does not shrink as task count
35 grows.
36 - PROCESS: separate memory per worker, needs pickling. Pays a large
37 fixed cost once, spawning workers and importing dependencies in
38 each -- which then amortizes as task count grows.
40 See [Parallelization](../getting_started/parallelization.html) for
41 measured trade-offs between the two.
42 """
44 THREAD = "thread"
45 PROCESS = "process"
48def _locate_worker(
49 args: tuple[PixelCoordinate, PixelCoordinate],
50 *,
51 reference_image: np.ndarray,
52 current_image: np.ndarray,
53 kernel_margin_width: int,
54 kernel_margin_height: int,
55 search_margin_width: int,
56 search_margin_height: int,
57 windowing: WindowingMethod | None,
58) -> PixelCoordinate:
59 """One point's worth of `translation.locate`, as a single positional
60 argument -- `Executor.map` (thread or process) always calls its
61 target positionally, one item per iterable, so a keyword-only
62 signature is not an option for the function being mapped over.
63 Module-level on purpose: `ProcessPoolExecutor` needs a real,
64 importable function to hand to spawned workers, not a closure.
65 """
66 reference_point, search_center = args
67 return translation.locate(
68 reference_image=reference_image,
69 current_image=current_image,
70 reference_point=reference_point,
71 search_center=search_center,
72 kernel_margin_width=kernel_margin_width,
73 kernel_margin_height=kernel_margin_height,
74 search_margin_width=search_margin_width,
75 search_margin_height=search_margin_height,
76 windowing=windowing,
77 )
80def _locate_subpixel_worker(
81 args: tuple[PixelCoordinate, PixelCoordinate],
82 *,
83 reference_image: np.ndarray,
84 current_image: np.ndarray,
85 kernel_margin_width: int,
86 kernel_margin_height: int,
87 search_margin_width: int,
88 search_margin_height: int,
89 windowing: WindowingMethod | None,
90 upsample_factor: int,
91) -> SubpixelCoordinate:
92 """One point's worth of `translation.locate_subpixel` -- see
93 `_locate_worker`'s own docstring for why this is a module-level
94 function taking a single positional argument, not a closure."""
95 reference_point, search_center = args
96 return translation.locate_subpixel(
97 reference_image=reference_image,
98 current_image=current_image,
99 reference_point=reference_point,
100 search_center=search_center,
101 kernel_margin_width=kernel_margin_width,
102 kernel_margin_height=kernel_margin_height,
103 search_margin_width=search_margin_width,
104 search_margin_height=search_margin_height,
105 windowing=windowing,
106 upsample_factor=upsample_factor,
107 )
110def _locate_warp_worker(
111 args: tuple[PixelCoordinate, PixelCoordinate],
112 *,
113 reference_image: np.ndarray,
114 current_image: np.ndarray,
115 reference_gradients: tuple[np.ndarray, np.ndarray],
116 current_coefficients: np.ndarray,
117 kernel_margin_width: int,
118 kernel_margin_height: int,
119 search_margin_width: int,
120 search_margin_height: int,
121 windowing: WindowingMethod | None,
122 max_iterations: int,
123 tolerance: float,
124) -> SubpixelCoordinate:
125 """One point's worth of `warp.locate`, reusing gradients and spline
126 coefficients computed once per batch -- see `_locate_worker`'s own
127 docstring for why this is a module-level function taking a single
128 positional argument, not a closure."""
129 reference_point, search_center = args
130 start = translation.locate(
131 reference_image=reference_image,
132 current_image=current_image,
133 reference_point=reference_point,
134 search_center=search_center,
135 kernel_margin_width=kernel_margin_width,
136 kernel_margin_height=kernel_margin_height,
137 search_margin_width=search_margin_width,
138 search_margin_height=search_margin_height,
139 windowing=windowing,
140 )
141 fitted, _znssd, _iterations = warp._refine(
142 reference_image=reference_image,
143 reference_gradients=reference_gradients,
144 current_coefficients=current_coefficients,
145 reference_point=reference_point,
146 initial_warp=warp._translation_warp(
147 reference_point=reference_point, start=start
148 ),
149 kernel_margin_width=kernel_margin_width,
150 kernel_margin_height=kernel_margin_height,
151 max_iterations=max_iterations,
152 tolerance=tolerance,
153 )
154 return SubpixelCoordinate(
155 x=reference_point.x + fitted[0, 2], y=reference_point.y + fitted[1, 2]
156 )
159def generate(
160 *,
161 origin: PixelCoordinate,
162 count_x: int,
163 count_y: int,
164 spacing_x: int,
165 spacing_y: int,
166) -> list[PixelCoordinate]:
167 """Generate a rectangular collection of points spanning x and y.
169 Points are returned in row-major order (top-left to bottom-right: all
170 of row 0 first, then row 1, and so on) -- point `i` sits at the same
171 index a later batch tracking call (see `locate`) returns its found
172 position at.
174 `count_x` and `count_y` need not be equal, and `spacing_x` and
175 `spacing_y` need not be equal -- this is a general rectangular
176 collection, not a square or uniformly-spaced one.
178 `count_x` or `count_y` of exactly `1` is allowed here -- a single
179 row or column of points is still a meaningful point-tracking grid on
180 its own. That's looser than [`elements`](#elements)'s own
181 requirement of `>= 2` along each axis, since forming even 1 finite
182 element needs 2 points per axis. The two functions check their own,
183 different preconditions independently -- `generate` doesn't know or
184 care about elements.
186 Args:
187 origin: Position of the top-left point, in the source image's
188 pixel reference frame.
189 count_x: Number of points along x. Must be >= 1.
190 count_y: Number of points along y. Must be >= 1.
191 spacing_x: Pixel spacing between adjacent points along x.
192 spacing_y: Pixel spacing between adjacent points along y.
194 Returns:
195 A list of `count_x * count_y` `PixelCoordinate`s, in row-major
196 order.
198 Raises:
199 ValueError: If `count_x` or `count_y` is less than 1.
200 """
201 if count_x < 1:
202 raise ValueError(f"count_x {count_x} must be >= 1")
203 if count_y < 1:
204 raise ValueError(f"count_y {count_y} must be >= 1")
206 return [
207 PixelCoordinate(x=origin.x + i * spacing_x, y=origin.y + j * spacing_y)
208 for j in range(count_y)
209 for i in range(count_x)
210 ]
213def elements(*, count_x: int, count_y: int) -> list[tuple[int, int, int, int]]:
214 """Q4 connectivity for the regular `count_x` x `count_y` lattice `generate` produces.
216 Each element is one of `generate`'s point grid's unit cells, its 4
217 corner nodes given as indices into that same points list -- so
218 `[points[i] for i in elements(...)[0]]` are one element's 4 corner
219 `PixelCoordinate`s, ready to hand to
220 [`dictk.element.gauss_point_green_lagrange_strains`](../api/dictk/element.html#gauss_point_green_lagrange_strains)
221 or
222 [`dictk.element.gauss_point_log_strains`](../api/dictk/element.html#gauss_point_log_strains).
224 Each 4-tuple is `(top_left, top_right, bottom_right, bottom_left)`
225 point indices -- the same $N_1$..$N_4$ corner order those functions'
226 `shape_functions` convention expects (see [Shape
227 Functions](../getting_started/finite_element_method.html#shape-functions)).
228 "Top"/"bottom" here use `generate`'s own image-pixel convention (y
229 increasing downward, origin at the top-left point) -- not the
230 math-style y-up convention `finite_element_method.md`'s figures use,
231 where the same 4 points would be called `N1`..`N4`'s "bottom-left,
232 bottom-right, top-right, top-left" instead. Only the labels differ
233 between the two: the actual point order (start at one corner, then
234 +x, then +x and +y together, then +y) is identical either way, and
235 that -- not which direction is called "up" -- is what makes it match
236 `shape_functions`' expected winding.
238 Args:
239 count_x: Number of points along x in the source `generate` call.
240 Must be >= 2 (at least 2 points make 1 element along x).
241 count_y: Number of points along y in the source `generate` call.
242 Must be >= 2.
244 Returns:
245 A list of `(count_x - 1) * (count_y - 1)` 4-tuples, in row-major
246 order (all of element row 0 first, then row 1, and so on) --
247 matching `generate`'s own point ordering.
249 Raises:
250 ValueError: If `count_x` or `count_y` is less than 2.
251 """
252 if count_x < 2:
253 raise ValueError(f"count_x {count_x} must be >= 2")
254 if count_y < 2:
255 raise ValueError(f"count_y {count_y} must be >= 2")
257 return [
258 (
259 j * count_x + i, # top_left
260 j * count_x + i + 1, # top_right
261 (j + 1) * count_x + i + 1, # bottom_right
262 (j + 1) * count_x + i, # bottom_left
263 )
264 for j in range(count_y - 1)
265 for i in range(count_x - 1)
266 ]
269def locate(
270 *,
271 reference_image: np.ndarray,
272 current_image: np.ndarray,
273 reference_points: Sequence[PixelCoordinate],
274 search_centers: Sequence[PixelCoordinate] | None = None,
275 kernel_margin_width: int,
276 kernel_margin_height: int,
277 search_margin_width: int,
278 search_margin_height: int,
279 windowing: WindowingMethod | None = None,
280 max_workers: int | None = None,
281 executor: Executor = Executor.THREAD,
282) -> list[PixelCoordinate]:
283 """Batch version of `dictk.translation.locate`: track many points at once.
285 Given a collection of `reference_points` (e.g. from `generate`), finds
286 each one's position in `current_image` by calling
287 [`dictk.translation.locate`](./translation.html#locate) once per point.
288 Returns a flat list, index-aligned with `reference_points` -- the same
289 row-major order `generate` produces, so point `i`'s found position is
290 at index `i` of the result.
292 Args:
293 reference_image: The reference (undeformed) 2D grayscale image.
294 current_image: The current (deformed) 2D grayscale image.
295 reference_points: Each point's fixed, known position, in
296 `reference_image`'s pixel reference frame.
297 search_centers: Where to center each point's search area, in
298 `current_image`'s pixel reference frame -- a guess of roughly
299 where each `reference_points` entry ended up. If `None`
300 (default), each point's own `reference_points` entry is used
301 as its own search center, the same reasonable default
302 `translation.locate` documents for the single-point case. If
303 given, must be the same length as `reference_points`.
304 kernel_margin_width: Half each kernel's width, in pixels. Must be
305 >= 1.
306 kernel_margin_height: Half each kernel's height, in pixels. Must
307 be >= 1.
308 search_margin_width: Half each search area's width, in pixels.
309 Must be greater than `kernel_margin_width`.
310 search_margin_height: Half each search area's height, in pixels.
311 Must be greater than `kernel_margin_height`.
312 windowing: Passed straight through to each per-point
313 [`dictk.translation.locate`](./translation.html#locate) call
314 -- see its own `windowing` parameter. Default `None` applies
315 no windowing to any point, matching this function's original
316 behavior exactly.
317 max_workers: If given, points are tracked concurrently across
318 this many workers instead of one at a time. Default `None`
319 stays sequential -- a plain loop, no pool, no overhead --
320 matching this function's original behavior exactly. See
321 [Parallelization](../getting_started/parallelization.html)
322 before setting this: at this book's own teaching scale
323 (small kernels and search areas), sequential outperforms
324 both pool types at any point count up to 1,000,000, measured.
325 Concurrency only pays for itself once each point's own
326 correlation is large enough, or point count is large enough
327 to amortize `Executor.PROCESS`'s fixed startup cost -- see
328 that page for the measured trade-offs.
329 executor: Which pool type `max_workers` runs on. Ignored if
330 `max_workers` is `None`. Default `Executor.THREAD`.
332 Returns:
333 Each point's location, in `current_image`'s pixel reference
334 frame, in the same order as `reference_points`.
336 Raises:
337 ValueError: If `search_centers` is given and its length doesn't
338 match `reference_points`, `max_workers` is given and less
339 than 1, or (from the underlying per-point call) if the
340 margin arguments are invalid.
341 """
342 if search_centers is None:
343 search_centers = reference_points
344 elif len(search_centers) != len(reference_points):
345 raise ValueError(
346 f"search_centers length ({len(search_centers)}) must match "
347 f"reference_points length ({len(reference_points)})"
348 )
349 if max_workers is not None and max_workers < 1:
350 raise ValueError(f"max_workers {max_workers} must be >= 1")
352 if max_workers is None:
353 return [
354 translation.locate(
355 reference_image=reference_image,
356 current_image=current_image,
357 reference_point=reference_point,
358 search_center=search_center,
359 kernel_margin_width=kernel_margin_width,
360 kernel_margin_height=kernel_margin_height,
361 search_margin_width=search_margin_width,
362 search_margin_height=search_margin_height,
363 windowing=windowing,
364 )
365 for reference_point, search_center in zip(reference_points, search_centers)
366 ]
368 worker = partial(
369 _locate_worker,
370 reference_image=reference_image,
371 current_image=current_image,
372 kernel_margin_width=kernel_margin_width,
373 kernel_margin_height=kernel_margin_height,
374 search_margin_width=search_margin_width,
375 search_margin_height=search_margin_height,
376 windowing=windowing,
377 )
378 executor_cls = (
379 ThreadPoolExecutor if executor is Executor.THREAD else ProcessPoolExecutor
380 )
381 with executor_cls(max_workers=max_workers) as pool:
382 return list(pool.map(worker, zip(reference_points, search_centers)))
385def locate_subpixel(
386 *,
387 reference_image: np.ndarray,
388 current_image: np.ndarray,
389 reference_points: Sequence[PixelCoordinate],
390 search_centers: Sequence[PixelCoordinate] | None = None,
391 kernel_margin_width: int,
392 kernel_margin_height: int,
393 search_margin_width: int,
394 search_margin_height: int,
395 windowing: WindowingMethod | None = None,
396 upsample_factor: int = 100,
397 max_workers: int | None = None,
398 executor: Executor = Executor.THREAD,
399) -> list[SubpixelCoordinate]:
400 """Batch version of `dictk.translation.locate_subpixel`: track many
401 points at once, with subpixel refinement.
403 Same structure as [`locate`](#locate) -- see its own docstring for
404 the full explanation of `search_centers`, `windowing`,
405 `max_workers`, and `executor`. The one behavioral difference beyond
406 the added `upsample_factor` parameter: each point's found position
407 is a [`SubpixelCoordinate`](../image.html#SubpixelCoordinate)
408 (generally fractional), not a `PixelCoordinate`, matching
409 `translation.locate_subpixel`'s own return type.
411 Args:
412 reference_image: The reference (undeformed) 2D grayscale image.
413 current_image: The current (deformed) 2D grayscale image.
414 reference_points: Each point's fixed, known position, in
415 `reference_image`'s pixel reference frame.
416 search_centers: Where to center each point's search area -- see
417 `locate`'s own docstring. If `None` (default), each point's
418 own `reference_points` entry is used.
419 kernel_margin_width: Half each kernel's width, in pixels. Must
420 be >= 1.
421 kernel_margin_height: Half each kernel's height, in pixels. Must
422 be >= 1.
423 search_margin_width: Half each search area's width, in pixels.
424 Must be greater than `kernel_margin_width`.
425 search_margin_height: Half each search area's height, in pixels.
426 Must be greater than `kernel_margin_height`.
427 windowing: Passed straight through to each per-point
428 [`dictk.translation.locate_subpixel`](../translation.html#locate_subpixel)
429 call. Default `None` applies no windowing.
430 upsample_factor: Passed straight through to each per-point
431 `locate_subpixel` call. Default `100`, matching that
432 function's own default. Must be >= 1.
433 max_workers: If given, points are tracked concurrently across
434 this many workers instead of one at a time -- see `locate`'s
435 own docstring for the measured trade-offs.
436 executor: Which pool type `max_workers` runs on. Ignored if
437 `max_workers` is `None`. Default `Executor.THREAD`.
439 Returns:
440 Each point's location, in `current_image`'s pixel reference
441 frame, generally fractional, in the same order as
442 `reference_points`.
444 Raises:
445 ValueError: If `search_centers` is given and its length doesn't
446 match `reference_points`, `max_workers` is given and less
447 than 1, or (from the underlying per-point call) if the
448 margin or `upsample_factor` arguments are invalid.
449 """
450 if search_centers is None:
451 search_centers = reference_points
452 elif len(search_centers) != len(reference_points):
453 raise ValueError(
454 f"search_centers length ({len(search_centers)}) must match "
455 f"reference_points length ({len(reference_points)})"
456 )
457 if max_workers is not None and max_workers < 1:
458 raise ValueError(f"max_workers {max_workers} must be >= 1")
460 if max_workers is None:
461 return [
462 translation.locate_subpixel(
463 reference_image=reference_image,
464 current_image=current_image,
465 reference_point=reference_point,
466 search_center=search_center,
467 kernel_margin_width=kernel_margin_width,
468 kernel_margin_height=kernel_margin_height,
469 search_margin_width=search_margin_width,
470 search_margin_height=search_margin_height,
471 windowing=windowing,
472 upsample_factor=upsample_factor,
473 )
474 for reference_point, search_center in zip(reference_points, search_centers)
475 ]
477 worker = partial(
478 _locate_subpixel_worker,
479 reference_image=reference_image,
480 current_image=current_image,
481 kernel_margin_width=kernel_margin_width,
482 kernel_margin_height=kernel_margin_height,
483 search_margin_width=search_margin_width,
484 search_margin_height=search_margin_height,
485 windowing=windowing,
486 upsample_factor=upsample_factor,
487 )
488 executor_cls = (
489 ThreadPoolExecutor if executor is Executor.THREAD else ProcessPoolExecutor
490 )
491 with executor_cls(max_workers=max_workers) as pool:
492 return list(pool.map(worker, zip(reference_points, search_centers)))
495def locate_warp(
496 *,
497 reference_image: np.ndarray,
498 current_image: np.ndarray,
499 reference_points: Sequence[PixelCoordinate],
500 search_centers: Sequence[PixelCoordinate] | None = None,
501 kernel_margin_width: int,
502 kernel_margin_height: int,
503 search_margin_width: int,
504 search_margin_height: int,
505 windowing: WindowingMethod | None = None,
506 max_iterations: int = 50,
507 tolerance: float = 1e-6,
508 max_workers: int | None = None,
509 executor: Executor = Executor.THREAD,
510) -> list[SubpixelCoordinate]:
511 """Batch version of `dictk.warp.locate`: track many points at once
512 by warping an affine kernel.
514 Same structure as [`locate_subpixel`](#locate_subpixel) -- see
515 [`locate`](#locate)'s own docstring for the full explanation of
516 `search_centers`, `windowing`, `max_workers`, and `executor`.
517 `reference_image`'s gradients and `current_image`'s spline
518 coefficients are computed once per call, then shared by every point.
520 Args:
521 reference_image: The reference (undeformed) 2D grayscale image.
522 current_image: The current (deformed) 2D grayscale image.
523 reference_points: Each point's fixed, known position, in
524 `reference_image`'s pixel reference frame.
525 search_centers: Where to center each point's whole-pixel search
526 -- see `locate`'s own docstring. If `None` (default), each
527 point's own `reference_points` entry is used.
528 kernel_margin_width: Half each kernel's width, in pixels, not
529 counting its center column. Must be >= 1.
530 kernel_margin_height: Half each kernel's height, in pixels, not
531 counting its center row. Must be >= 1.
532 search_margin_width: Half each whole-pixel search area's width,
533 in pixels. Must be greater than `kernel_margin_width`.
534 search_margin_height: Half each whole-pixel search area's
535 height, in pixels. Must be greater than
536 `kernel_margin_height`.
537 windowing: Passed straight through to each point's whole-pixel
538 stage. Default `None` applies no windowing.
539 max_iterations: The most IC-GN updates per point. Default `50`.
540 Must be >= 1.
541 tolerance: IC-GN stops once every warp parameter's update falls
542 below this. Default `1e-6`. Must be > 0.
543 max_workers: If given, points are tracked concurrently across
544 this many workers instead of one at a time.
545 executor: Which pool type `max_workers` runs on. Ignored if
546 `max_workers` is `None`. Default `Executor.THREAD`.
548 Returns:
549 Each point's location, in `current_image`'s pixel reference
550 frame, generally fractional, in the same order as
551 `reference_points`.
553 Raises:
554 ValueError: If `search_centers` is given but its length doesn't
555 match `reference_points`, `max_workers` is less than 1,
556 `max_iterations` is less than 1, `tolerance` is not
557 positive, or any margin fails `dictk.translation.locate`'s
558 own checks.
559 """
560 if search_centers is None:
561 search_centers = reference_points
562 elif len(search_centers) != len(reference_points):
563 raise ValueError(
564 f"search_centers length ({len(search_centers)}) must match "
565 f"reference_points length ({len(reference_points)})"
566 )
567 if max_workers is not None and max_workers < 1:
568 raise ValueError(f"max_workers {max_workers} must be >= 1")
569 if max_iterations < 1:
570 raise ValueError(f"max_iterations {max_iterations} must be >= 1")
571 if tolerance <= 0:
572 raise ValueError(f"tolerance {tolerance} must be > 0")
574 worker = partial(
575 _locate_warp_worker,
576 reference_image=reference_image,
577 current_image=current_image,
578 reference_gradients=warp._gradients(reference_image),
579 current_coefficients=warp._spline_coefficients(current_image),
580 kernel_margin_width=kernel_margin_width,
581 kernel_margin_height=kernel_margin_height,
582 search_margin_width=search_margin_width,
583 search_margin_height=search_margin_height,
584 windowing=windowing,
585 max_iterations=max_iterations,
586 tolerance=tolerance,
587 )
588 if max_workers is None:
589 return [worker(pair) for pair in zip(reference_points, search_centers)]
591 executor_cls = (
592 ThreadPoolExecutor if executor is Executor.THREAD else ProcessPoolExecutor
593 )
594 with executor_cls(max_workers=max_workers) as pool:
595 return list(pool.map(worker, zip(reference_points, search_centers)))
598def locate_warp_propagated(
599 *,
600 reference_image: np.ndarray,
601 current_image: np.ndarray,
602 reference_points: Sequence[PixelCoordinate],
603 kernel_margin_width: int,
604 kernel_margin_height: int,
605 search_margin_width: int,
606 search_margin_height: int,
607 start_gradient_limit: float = 0.3,
608 start_gradient_step: float = 0.1,
609 znssd_limit: float = 0.3,
610 max_iterations: int = 50,
611 tolerance: float = 1e-6,
612) -> tuple[list[SubpixelCoordinate], np.ndarray]:
613 """Track many points by warping an affine kernel, seeding each point
614 from an already-fitted neighbor instead of a rigid whole-pixel match.
616 [`locate_warp`](#locate_warp) seeds every point's IC-GN fit from
617 `dictk.translation.locate`'s rigid, whole-pixel match. Under a large
618 stretch or shear, that match can land on the wrong speckles.
620 This function searches for a deformed seed at one point only: the
621 point nearest the center of `reference_image`. It tries every
622 displacement gradient `(u_x, u_y, v_x, v_y)` on a grid from
623 `-start_gradient_limit` to `start_gradient_limit` in steps of
624 `start_gradient_step`. For each one, it deforms the reference
625 kernel and finds its whole-pixel translation by phase correlation.
626 IC-GN refines the best-scoring trials, and the fit with the lowest
627 ZNSSD becomes the first point's warp.
629 A fit is accepted when its ZNSSD is at most `znssd_limit`. Only
630 accepted fits seed other points. The remaining points are visited
631 best first: the accepted point with the lowest ZNSSD seeds each of
632 its `NEIGHBOR_COUNT` nearest points not yet accepted. Each starts
633 from that neighbor's warp: the same displacement gradients, and a
634 translation extrapolated across the offset between the two points.
635 A point that fails `znssd_limit` tries again when another accepted
636 neighbor reaches it. See [Seed
637 Propagation](../getting_started/seed_propagation.html).
639 If the first point's fit fails `znssd_limit`, the start search
640 moves to the next untried point nearest the image center, skipping
641 points within one kernel width of a failed start, since those share
642 most of its kernel. Points
643 no accepted fit can reach get their own start search the same way.
644 The search stops after `START_ATTEMPT_LIMIT` failed starts in a
645 row. Start points whose search area leaves `current_image` are
646 skipped.
648 Only start points use the search area, centered on each one's own
649 reference position. Every other point skips the
650 whole-pixel search, so the search margins do not limit how far it
651 can move. The fit runs one point at a time, because each seed
652 depends on an earlier fit.
654 Args:
655 reference_image: The reference (undeformed) 2D grayscale image.
656 current_image: The current (deformed) 2D grayscale image.
657 reference_points: Each point's fixed, known position, in
658 `reference_image`'s pixel reference frame. Each point's
659 kernel must lie inside `reference_image`.
660 kernel_margin_width: Half each kernel's width, in pixels, not
661 counting its center column. Must be >= 1.
662 kernel_margin_height: Half each kernel's height, in pixels, not
663 counting its center row. Must be >= 1.
664 search_margin_width: Half the whole-pixel search area's width,
665 in pixels, for start points. Must be greater than
666 `kernel_margin_width`.
667 search_margin_height: Half the whole-pixel search area's
668 height, in pixels, for start points. Must be greater than
669 `kernel_margin_height`.
670 start_gradient_limit: The largest displacement gradient the
671 start search tries, in each of its four terms. Default
672 `0.3`. Must be > 0.
673 start_gradient_step: The spacing between tried gradients.
674 Default `0.1`, which tries 7 values per term and 2401 in
675 all. Must be > 0.
676 znssd_limit: The largest ZNSSD a fit may have and still seed
677 other points. Default `0.3`. Must be > 0.
678 max_iterations: The most IC-GN updates per point. Default `50`.
679 Must be >= 1.
680 tolerance: IC-GN stops once every warp parameter's update falls
681 below this. Default `1e-6`. Must be > 0.
683 Returns:
684 Two items, each in the same order as `reference_points`:
686 - Each point's location, in `current_image`'s pixel reference
687 frame, generally fractional. It comes from the point's
688 lowest-ZNSSD fit, even when that fit failed `znssd_limit`.
689 A point never fitted at all gets `nan` for `x` and `y`.
690 - Each point's ZNSSD for that fit, from 0 for a perfect match
691 to 4, or `inf` for a point never fitted. Values above
692 `znssd_limit` mark positions not to trust.
694 Raises:
695 ValueError: If `reference_points` is empty, any point's kernel
696 leaves `reference_image`, a margin is too small,
697 `start_gradient_limit`, `start_gradient_step`, or
698 `znssd_limit` is not positive, `max_iterations` is less
699 than 1, or `tolerance` is not positive.
700 """
701 if len(reference_points) == 0:
702 raise ValueError("reference_points must not be empty")
703 if kernel_margin_width < 1:
704 raise ValueError(f"kernel_margin_width {kernel_margin_width} must be >= 1")
705 if kernel_margin_height < 1:
706 raise ValueError(f"kernel_margin_height {kernel_margin_height} must be >= 1")
707 if search_margin_width <= kernel_margin_width:
708 raise ValueError(
709 f"search_margin_width {search_margin_width} must be greater than "
710 f"kernel_margin_width {kernel_margin_width}"
711 )
712 if search_margin_height <= kernel_margin_height:
713 raise ValueError(
714 f"search_margin_height {search_margin_height} must be greater than "
715 f"kernel_margin_height {kernel_margin_height}"
716 )
717 if start_gradient_limit <= 0:
718 raise ValueError(f"start_gradient_limit {start_gradient_limit} must be > 0")
719 if start_gradient_step <= 0:
720 raise ValueError(f"start_gradient_step {start_gradient_step} must be > 0")
721 if znssd_limit <= 0:
722 raise ValueError(f"znssd_limit {znssd_limit} must be > 0")
723 if max_iterations < 1:
724 raise ValueError(f"max_iterations {max_iterations} must be >= 1")
725 if tolerance <= 0:
726 raise ValueError(f"tolerance {tolerance} must be > 0")
727 height, width = reference_image.shape
728 for point in reference_points:
729 if not (
730 kernel_margin_width <= point.x < width - kernel_margin_width
731 and kernel_margin_height <= point.y < height - kernel_margin_height
732 ):
733 raise ValueError(f"kernel around {point} leaves reference_image")
735 positions = np.array([[p.x, p.y] for p in reference_points], dtype=np.float64)
736 neighbor_count = min(NEIGHBOR_COUNT + 1, len(reference_points))
737 _, neighbors = cKDTree(positions).query(positions, k=neighbor_count)
738 neighbors = np.asarray(neighbors).reshape(len(reference_points), -1)
739 center = np.array([(width - 1) / 2, (height - 1) / 2])
740 by_distance = np.argsort(np.linalg.norm(positions - center, axis=1))
742 reference_gradients = warp._gradients(reference_image)
743 reference_coefficients = warp._spline_coefficients(reference_image)
744 current_coefficients = warp._spline_coefficients(current_image)
746 def refined(index, initial_warp):
747 return warp._refine(
748 reference_image=reference_image,
749 reference_gradients=reference_gradients,
750 current_coefficients=current_coefficients,
751 reference_point=reference_points[index],
752 initial_warp=initial_warp,
753 kernel_margin_width=kernel_margin_width,
754 kernel_margin_height=kernel_margin_height,
755 max_iterations=max_iterations,
756 tolerance=tolerance,
757 )
759 count = len(reference_points)
760 warps: list[np.ndarray | None] = [None] * count
761 znssds = np.full(count, np.inf)
762 accepted = np.zeros(count, dtype=bool)
763 tried = np.zeros(count, dtype=bool)
764 current_height, current_width = current_image.shape
766 def record(index, fit):
767 fitted, znssd, _iterations = fit
768 tried[index] = True
769 if znssd < znssds[index]:
770 warps[index], znssds[index] = fitted, znssd
771 if znssd <= znssd_limit:
772 accepted[index] = True
773 return accepted[index]
775 def propagate(first):
776 heap = [(znssds[first], first)]
777 while heap:
778 _, parent = heapq.heappop(heap)
779 parent_warp = warps[parent]
780 for child in neighbors[parent]:
781 if accepted[child]:
782 continue
783 offset = positions[child] - positions[parent]
784 initial_warp = parent_warp.copy()
785 initial_warp[:2, 2] += (parent_warp[:2, :2] - np.eye(2)) @ offset
786 if record(child, refined(child, initial_warp)):
787 heapq.heappush(heap, (znssds[child], child))
789 failed_starts = 0
790 failed_start_positions: list[np.ndarray] = []
791 kernel_width = 2 * max(kernel_margin_width, kernel_margin_height) + 1
792 for first in by_distance:
793 if tried[first] or failed_starts >= START_ATTEMPT_LIMIT:
794 continue
795 if any(
796 np.linalg.norm(positions[first] - failed) < kernel_width
797 for failed in failed_start_positions
798 ):
799 continue
800 point = reference_points[first]
801 if not (
802 search_margin_width <= point.x <= current_width - search_margin_width
803 and search_margin_height <= point.y <= current_height - search_margin_height
804 ):
805 continue
806 seeds = warp._shape_seeds(
807 reference_coefficients=reference_coefficients,
808 current_image=current_image,
809 current_coefficients=current_coefficients,
810 reference_point=point,
811 kernel_margin_width=kernel_margin_width,
812 kernel_margin_height=kernel_margin_height,
813 search_margin_width=search_margin_width,
814 search_margin_height=search_margin_height,
815 gradient_limit=start_gradient_limit,
816 gradient_step=start_gradient_step,
817 seed_count=START_SEED_COUNT,
818 )
819 best = min((refined(first, seed) for seed in seeds), key=lambda fit: fit[1])
820 if record(first, best):
821 failed_starts = 0
822 propagate(first)
823 else:
824 failed_starts += 1
825 failed_start_positions.append(positions[first])
827 positions_found = [
828 SubpixelCoordinate(x=p.x + w[0, 2], y=p.y + w[1, 2])
829 if w is not None
830 else SubpixelCoordinate(x=np.nan, y=np.nan)
831 for p, w in zip(reference_points, warps)
832 ]
833 return positions_found, znssds