setiastrosuitepro 1.6.10__py3-none-any.whl → 1.7.0.post2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (51) hide show
  1. setiastro/images/colorwheel.svg +97 -0
  2. setiastro/images/narrowbandnormalization.png +0 -0
  3. setiastro/images/planetarystacker.png +0 -0
  4. setiastro/saspro/__main__.py +1 -1
  5. setiastro/saspro/_generated/build_info.py +2 -2
  6. setiastro/saspro/aberration_ai.py +49 -11
  7. setiastro/saspro/aberration_ai_preset.py +29 -3
  8. setiastro/saspro/backgroundneutral.py +73 -33
  9. setiastro/saspro/blink_comparator_pro.py +116 -71
  10. setiastro/saspro/convo.py +9 -6
  11. setiastro/saspro/curve_editor_pro.py +72 -22
  12. setiastro/saspro/curves_preset.py +249 -47
  13. setiastro/saspro/doc_manager.py +178 -11
  14. setiastro/saspro/gui/main_window.py +305 -66
  15. setiastro/saspro/gui/mixins/dock_mixin.py +245 -24
  16. setiastro/saspro/gui/mixins/file_mixin.py +35 -16
  17. setiastro/saspro/gui/mixins/menu_mixin.py +32 -1
  18. setiastro/saspro/gui/mixins/toolbar_mixin.py +135 -11
  19. setiastro/saspro/histogram.py +179 -7
  20. setiastro/saspro/imageops/narrowband_normalization.py +816 -0
  21. setiastro/saspro/imageops/serloader.py +972 -0
  22. setiastro/saspro/imageops/starbasedwhitebalance.py +23 -52
  23. setiastro/saspro/imageops/stretch.py +66 -15
  24. setiastro/saspro/legacy/numba_utils.py +25 -48
  25. setiastro/saspro/live_stacking.py +24 -4
  26. setiastro/saspro/multiscale_decomp.py +30 -17
  27. setiastro/saspro/narrowband_normalization.py +1618 -0
  28. setiastro/saspro/numba_utils.py +0 -55
  29. setiastro/saspro/ops/script_editor.py +5 -0
  30. setiastro/saspro/ops/scripts.py +119 -0
  31. setiastro/saspro/remove_green.py +1 -1
  32. setiastro/saspro/resources.py +4 -0
  33. setiastro/saspro/ser_stack_config.py +74 -0
  34. setiastro/saspro/ser_stacker.py +2310 -0
  35. setiastro/saspro/ser_stacker_dialog.py +1500 -0
  36. setiastro/saspro/ser_tracking.py +206 -0
  37. setiastro/saspro/serviewer.py +1258 -0
  38. setiastro/saspro/sfcc.py +602 -214
  39. setiastro/saspro/shortcuts.py +35 -16
  40. setiastro/saspro/stacking_suite.py +332 -87
  41. setiastro/saspro/star_alignment.py +243 -122
  42. setiastro/saspro/stat_stretch.py +220 -31
  43. setiastro/saspro/subwindow.py +2 -4
  44. setiastro/saspro/whitebalance.py +24 -0
  45. setiastro/saspro/widgets/resource_monitor.py +122 -74
  46. {setiastrosuitepro-1.6.10.dist-info → setiastrosuitepro-1.7.0.post2.dist-info}/METADATA +2 -2
  47. {setiastrosuitepro-1.6.10.dist-info → setiastrosuitepro-1.7.0.post2.dist-info}/RECORD +51 -40
  48. {setiastrosuitepro-1.6.10.dist-info → setiastrosuitepro-1.7.0.post2.dist-info}/WHEEL +0 -0
  49. {setiastrosuitepro-1.6.10.dist-info → setiastrosuitepro-1.7.0.post2.dist-info}/entry_points.txt +0 -0
  50. {setiastrosuitepro-1.6.10.dist-info → setiastrosuitepro-1.7.0.post2.dist-info}/licenses/LICENSE +0 -0
  51. {setiastrosuitepro-1.6.10.dist-info → setiastrosuitepro-1.7.0.post2.dist-info}/licenses/license.txt +0 -0
@@ -0,0 +1,2310 @@
1
+ # src/setiastro/saspro/ser_stacker.py
2
+ from __future__ import annotations
3
+ import os
4
+ import threading
5
+
6
+ from concurrent.futures import ThreadPoolExecutor, as_completed
7
+
8
+ from dataclasses import dataclass
9
+ from typing import Optional, Tuple, List, Dict, Any
10
+
11
+ import numpy as np
12
+
13
+ import cv2
14
+ cv2.setNumThreads(1)
15
+
16
+ from setiastro.saspro.imageops.serloader import SERReader
17
+ from setiastro.saspro.ser_stack_config import SERStackConfig
18
+ from setiastro.saspro.ser_tracking import PlanetaryTracker, SurfaceTracker, _to_mono01
19
+ from setiastro.saspro.imageops.serloader import open_planetary_source, PlanetaryFrameSource
20
+
21
+ _BAYER_TO_CV2 = {
22
+ "RGGB": cv2.COLOR_BayerRG2RGB,
23
+ "BGGR": cv2.COLOR_BayerBG2RGB,
24
+ "GRBG": cv2.COLOR_BayerGR2RGB,
25
+ "GBRG": cv2.COLOR_BayerGB2RGB,
26
+ }
27
+
28
+ def _cfg_bayer_pattern(cfg) -> str | None:
29
+ # cfg.bayer_pattern might be missing in older saved projects; be defensive
30
+ return getattr(cfg, "bayer_pattern", None)
31
+
32
+
33
+ def _get_frame(src, idx: int, *, roi, debayer: bool, to_float01: bool, force_rgb: bool, bayer_pattern: str | None):
34
+ """
35
+ Drop-in wrapper:
36
+ - passes cfg.bayer_pattern down to sources that support it
37
+ - stays compatible with sources whose get_frame() doesn't accept bayer_pattern yet
38
+ """
39
+ try:
40
+ return src.get_frame(
41
+ int(idx),
42
+ roi=roi,
43
+ debayer=debayer,
44
+ to_float01=to_float01,
45
+ force_rgb=force_rgb,
46
+ bayer_pattern=bayer_pattern,
47
+ )
48
+ except TypeError:
49
+ # Back-compat: older PlanetaryFrameSource implementations
50
+ return src.get_frame(
51
+ int(idx),
52
+ roi=roi,
53
+ debayer=debayer,
54
+ to_float01=to_float01,
55
+ force_rgb=force_rgb,
56
+ )
57
+
58
+
59
+ @dataclass
60
+ class AnalyzeResult:
61
+ frames_total: int
62
+ roi_used: Optional[Tuple[int, int, int, int]]
63
+ track_mode: str
64
+ quality: np.ndarray # (N,) float32 higher=better
65
+ dx: np.ndarray # (N,) float32
66
+ dy: np.ndarray # (N,) float32
67
+ conf: np.ndarray # (N,) float32 0..1 (final conf used by stacking)
68
+ order: np.ndarray # (N,) int indices sorted by quality desc
69
+ ref_mode: str # "best_frame" | "best_stack"
70
+ ref_count: int
71
+ ref_image: np.ndarray # float32 [0..1], ROI-sized
72
+ ap_centers: Optional[np.ndarray] = None # (M,2) int32 in ROI coords
73
+ ap_size: int = 64
74
+ ap_multiscale: bool = False
75
+
76
+ # ✅ NEW: surface anchor confidence (coarse tracker)
77
+ coarse_conf: Optional[np.ndarray] = None # (N,) float32 0..1
78
+
79
+
80
+ @dataclass
81
+ class FrameEval:
82
+ idx: int
83
+ score: float
84
+ dx: float
85
+ dy: float
86
+ conf: float
87
+
88
+ def _print_surface_debug(
89
+ *,
90
+ dx: np.ndarray,
91
+ dy: np.ndarray,
92
+ conf: np.ndarray,
93
+ coarse_conf: np.ndarray | None,
94
+ floor: float = 0.05,
95
+ prefix: str = "[SER][Surface]"
96
+ ) -> None:
97
+ try:
98
+ dx = np.asarray(dx, dtype=np.float32)
99
+ dy = np.asarray(dy, dtype=np.float32)
100
+ conf = np.asarray(conf, dtype=np.float32)
101
+
102
+ dx_min = float(np.min(dx)) if dx.size else 0.0
103
+ dx_max = float(np.max(dx)) if dx.size else 0.0
104
+ dy_min = float(np.min(dy)) if dy.size else 0.0
105
+ dy_max = float(np.max(dy)) if dy.size else 0.0
106
+
107
+ conf_mean = float(np.mean(conf)) if conf.size else 0.0
108
+ conf_min = float(np.min(conf)) if conf.size else 0.0
109
+
110
+ msg = (
111
+ f"{prefix} dx[min,max]=({dx_min:.2f},{dx_max:.2f}) "
112
+ f"dy[min,max]=({dy_min:.2f},{dy_max:.2f}) "
113
+ f"conf[mean,min]=({conf_mean:.3f},{conf_min:.3f})"
114
+ )
115
+
116
+ if coarse_conf is not None:
117
+ cc = np.asarray(coarse_conf, dtype=np.float32)
118
+ cc_mean = float(np.mean(cc)) if cc.size else 0.0
119
+ cc_min = float(np.min(cc)) if cc.size else 0.0
120
+ cc_bad = float(np.mean(cc < 0.2)) if cc.size else 0.0
121
+ msg += f" coarse_conf[mean,min]=({cc_mean:.3f},{cc_min:.3f}) frac<0.2={cc_bad:.2%}"
122
+
123
+ if conf_mean <= floor + 1e-6:
124
+ msg += f" ⚠ conf.mean near floor ({floor}); alignment likely failing"
125
+ print(msg)
126
+ except Exception as e:
127
+ print(f"{prefix} debug print failed: {e}")
128
+
129
+
130
+ def _clamp_roi_in_bounds(roi: Tuple[int, int, int, int], w: int, h: int) -> Tuple[int, int, int, int]:
131
+ x, y, rw, rh = [int(v) for v in roi]
132
+ x = max(0, min(w - 1, x))
133
+ y = max(0, min(h - 1, y))
134
+ rw = max(1, min(w - x, rw))
135
+ rh = max(1, min(h - y, rh))
136
+ return x, y, rw, rh
137
+
138
+ def _grad_img(m: np.ndarray) -> np.ndarray:
139
+ """Simple, robust edge image for SSD refine."""
140
+ m = m.astype(np.float32, copy=False)
141
+ if cv2 is None:
142
+ # fallback: finite differences
143
+ gx = np.zeros_like(m); gx[:, 1:] = m[:, 1:] - m[:, :-1]
144
+ gy = np.zeros_like(m); gy[1:, :] = m[1:, :] - m[:-1, :]
145
+ g = np.abs(gx) + np.abs(gy)
146
+ g -= float(g.mean())
147
+ s = float(g.std()) + 1e-6
148
+ return g / s
149
+
150
+ gx = cv2.Sobel(m, cv2.CV_32F, 1, 0, ksize=3)
151
+ gy = cv2.Sobel(m, cv2.CV_32F, 0, 1, ksize=3)
152
+ g = cv2.magnitude(gx, gy)
153
+ g -= float(g.mean())
154
+ s = float(g.std()) + 1e-6
155
+ return (g / s).astype(np.float32, copy=False)
156
+
157
+ def _ssd_prepare_ref(ref_m: np.ndarray, crop: float = 0.80):
158
+ """
159
+ Precompute reference gradient + crop window once.
160
+
161
+ Returns:
162
+ rg : full reference gradient image (float32)
163
+ rgc : cropped view of rg
164
+ sl : (y0,y1,x0,x1) crop slices
165
+ """
166
+ ref_m = ref_m.astype(np.float32, copy=False)
167
+ rg = _grad_img(ref_m) # compute ONCE
168
+
169
+ H, W = rg.shape[:2]
170
+ cfx = max(8, int(W * (1.0 - float(crop)) * 0.5))
171
+ cfy = max(8, int(H * (1.0 - float(crop)) * 0.5))
172
+ x0, x1 = cfx, W - cfx
173
+ y0, y1 = cfy, H - cfy
174
+
175
+ rgc = rg[y0:y1, x0:x1] # view
176
+ return rg, rgc, (y0, y1, x0, x1)
177
+
178
+ def _subpixel_quadratic_1d(vm: float, v0: float, vp: float) -> float:
179
+ """
180
+ Given SSD at (-1,0,+1): (vm, v0, vp), return vertex offset in [-0.5,0.5]-ish.
181
+ Works for minimizing SSD.
182
+ """
183
+ denom = (vm - 2.0 * v0 + vp)
184
+ if abs(denom) < 1e-12:
185
+ return 0.0
186
+ # vertex of parabola fit through -1,0,+1
187
+ t = 0.5 * (vm - vp) / denom
188
+ return float(np.clip(t, -0.75, 0.75))
189
+
190
+
191
+ def _ssd_confidence_prepared(
192
+ rgc: np.ndarray,
193
+ cgc0: np.ndarray,
194
+ dx_i: int,
195
+ dy_i: int,
196
+ ) -> float:
197
+ """
198
+ Compute SSD between rgc and cgc0 shifted by (dx_i,dy_i) using slicing overlap.
199
+ Returns SSD (lower is better).
200
+
201
+ NOTE: This is integer-only and extremely fast (no warps).
202
+ """
203
+ H, W = rgc.shape[:2]
204
+
205
+ # Overlap slices for rgc and shifted cgc0
206
+ x0r = max(0, dx_i)
207
+ x1r = min(W, W + dx_i)
208
+ y0r = max(0, dy_i)
209
+ y1r = min(H, H + dy_i)
210
+
211
+ x0c = max(0, -dx_i)
212
+ x1c = min(W, W - dx_i)
213
+ y0c = max(0, -dy_i)
214
+ y1c = min(H, H - dy_i)
215
+
216
+ rr = rgc[y0r:y1r, x0r:x1r]
217
+ cc = cgc0[y0c:y1c, x0c:x1c]
218
+
219
+ d = rr - cc
220
+ return float(np.mean(d * d))
221
+
222
+
223
+ def _ssd_confidence(
224
+ ref_m: np.ndarray,
225
+ cur_m: np.ndarray,
226
+ dx: float,
227
+ dy: float,
228
+ *,
229
+ crop: float = 0.80,
230
+ ) -> float:
231
+ """
232
+ Original API: confidence from gradient SSD, higher=better (0..1).
233
+
234
+ Optimized:
235
+ - computes ref grad once per call (still OK if used standalone)
236
+ - uses one warp for (dx,dy)
237
+ - no extra work beyond necessary
238
+
239
+ For iterative search, use _refine_shift_ssd() which avoids redoing work.
240
+ """
241
+ ref_m = ref_m.astype(np.float32, copy=False)
242
+ cur_m = cur_m.astype(np.float32, copy=False)
243
+
244
+ # shift current by the proposed shift
245
+ cur_s = _shift_image(cur_m, float(dx), float(dy))
246
+
247
+ rg, rgc, sl = _ssd_prepare_ref(ref_m, crop=crop)
248
+ y0, y1, x0, x1 = sl
249
+
250
+ cg = _grad_img(cur_s)
251
+ cgc = cg[y0:y1, x0:x1]
252
+
253
+ d = rgc - cgc
254
+ ssd = float(np.mean(d * d))
255
+
256
+ scale = 0.002
257
+ conf = float(np.exp(-ssd / max(1e-12, scale)))
258
+ return float(np.clip(conf, 0.0, 1.0))
259
+
260
+
261
+ def _refine_shift_ssd(
262
+ ref_m: np.ndarray,
263
+ cur_m: np.ndarray,
264
+ dx0: float,
265
+ dy0: float,
266
+ *,
267
+ radius: int = 10,
268
+ crop: float = 0.80,
269
+ bruteforce: bool = False,
270
+ max_steps: int | None = None,
271
+ ) -> tuple[float, float, float]:
272
+ """
273
+ Returns (dx_refine, dy_refine, conf) where you ADD refine to (dx0,dy0).
274
+
275
+ CPU-optimized:
276
+ - precompute ref gradient crop once
277
+ - apply (dx0,dy0) shift ONCE
278
+ - compute gradient ONCE for shifted cur
279
+ - evaluate integer candidates via slicing overlap SSD (no warps)
280
+
281
+ If bruteforce=True, does full window scan in [-r,r]^2 (fast).
282
+ Otherwise does 8-neighbor hill-climb over integer offsets (very fast).
283
+
284
+ Optional subpixel polish:
285
+ - after choosing best integer (best_dx,best_dy), do a tiny separable quadratic
286
+ fit along x and y using SSD at +/-1 around the best integer.
287
+ - does NOT require any new gradients/warps (just 4 extra SSD evals).
288
+ """
289
+ r = int(max(0, radius))
290
+ if r == 0:
291
+ # nothing to do; just compute confidence at dx0/dy0
292
+ c = _ssd_confidence(ref_m, cur_m, dx0, dy0, crop=crop)
293
+ return 0.0, 0.0, float(c)
294
+
295
+ # Prepare ref grad crop ONCE
296
+ _, rgc, sl = _ssd_prepare_ref(ref_m, crop=crop)
297
+ y0, y1, x0, x1 = sl
298
+
299
+ # Shift cur by the current estimate ONCE, then gradient ONCE
300
+ cur_m = cur_m.astype(np.float32, copy=False)
301
+ cur0 = _shift_image(cur_m, float(dx0), float(dy0))
302
+ cg0 = _grad_img(cur0)
303
+ cgc0 = cg0[y0:y1, x0:x1]
304
+
305
+ # Helper: parabola vertex for minimizing SSD, using (-1,0,+1) samples
306
+ def _quad_min_offset(vm: float, v0: float, vp: float) -> float:
307
+ denom = (vm - 2.0 * v0 + vp)
308
+ if abs(denom) < 1e-12:
309
+ return 0.0
310
+ t = 0.5 * (vm - vp) / denom
311
+ return float(np.clip(t, -0.75, 0.75))
312
+
313
+ if bruteforce:
314
+ # NOTE: your bruteforce path currently includes a subpixel step already.
315
+ # If you want to keep using that exact implementation, just call it:
316
+ dxr, dyr, conf = _refine_shift_ssd_bruteforce(ref_m, cur_m, dx0, dy0, radius=r, crop=crop)
317
+ return float(dxr), float(dyr), float(conf)
318
+
319
+ # Hill-climb in integer space minimizing SSD
320
+ if max_steps is None:
321
+ max_steps = max(1, min(r, 6)) # small cap helps speed; tune if you want
322
+
323
+ best_dx = 0
324
+ best_dy = 0
325
+ best_ssd = _ssd_confidence_prepared(rgc, cgc0, 0, 0)
326
+
327
+ neigh = ((-1,0),(1,0),(0,-1),(0,1),(-1,-1),(-1,1),(1,-1),(1,1))
328
+
329
+ for _ in range(int(max_steps)):
330
+ improved = False
331
+ for sx, sy in neigh:
332
+ cand_dx = best_dx + sx
333
+ cand_dy = best_dy + sy
334
+ if abs(cand_dx) > r or abs(cand_dy) > r:
335
+ continue
336
+
337
+ ssd = _ssd_confidence_prepared(rgc, cgc0, cand_dx, cand_dy)
338
+ if ssd < best_ssd:
339
+ best_ssd = ssd
340
+ best_dx = cand_dx
341
+ best_dy = cand_dy
342
+ improved = True
343
+
344
+ if not improved:
345
+ break
346
+
347
+ # ---- subpixel quadratic polish around best integer (cheap) ----
348
+ # Uses SSD at +/-1 around best integer in X and Y (separable).
349
+ dx_sub = 0.0
350
+ dy_sub = 0.0
351
+ if r >= 1:
352
+ # X samples at (best_dx-1, best_dy), (best_dx, best_dy), (best_dx+1, best_dy)
353
+ if abs(best_dx - 1) <= r:
354
+ s_xm = _ssd_confidence_prepared(rgc, cgc0, best_dx - 1, best_dy)
355
+ else:
356
+ s_xm = best_ssd
357
+ s_x0 = best_ssd
358
+ if abs(best_dx + 1) <= r:
359
+ s_xp = _ssd_confidence_prepared(rgc, cgc0, best_dx + 1, best_dy)
360
+ else:
361
+ s_xp = best_ssd
362
+ dx_sub = _quad_min_offset(s_xm, s_x0, s_xp)
363
+
364
+ # Y samples at (best_dx, best_dy-1), (best_dx, best_dy), (best_dx, best_dy+1)
365
+ if abs(best_dy - 1) <= r:
366
+ s_ym = _ssd_confidence_prepared(rgc, cgc0, best_dx, best_dy - 1)
367
+ else:
368
+ s_ym = best_ssd
369
+ s_y0 = best_ssd
370
+ if abs(best_dy + 1) <= r:
371
+ s_yp = _ssd_confidence_prepared(rgc, cgc0, best_dx, best_dy + 1)
372
+ else:
373
+ s_yp = best_ssd
374
+ dy_sub = _quad_min_offset(s_ym, s_y0, s_yp)
375
+
376
+ best_dx_f = float(best_dx) + float(dx_sub)
377
+ best_dy_f = float(best_dy) + float(dy_sub)
378
+
379
+ # Confidence: keep based on best *integer* SSD (no subpixel warp needed)
380
+ scale = 0.002
381
+ conf = float(np.exp(-best_ssd / max(1e-12, scale)))
382
+ conf = float(np.clip(conf, 0.0, 1.0))
383
+
384
+ return float(best_dx_f), float(best_dy_f), float(conf)
385
+
386
+
387
+
388
+ def _refine_shift_ssd_bruteforce(
389
+ ref_m: np.ndarray,
390
+ cur_m: np.ndarray,
391
+ dx0: float,
392
+ dy0: float,
393
+ *,
394
+ radius: int = 2,
395
+ crop: float = 0.80,
396
+ ) -> tuple[float, float, float]:
397
+ """
398
+ Full brute-force scan in [-radius,+radius]^2, but optimized:
399
+ - shift by (dx0,dy0) ONCE
400
+ - compute gradients ONCE
401
+ - evaluate candidates via slicing overlap SSD (no warps)
402
+ - keep your separable quadratic subpixel fit
403
+ """
404
+ ref_m = ref_m.astype(np.float32, copy=False)
405
+ cur_m = cur_m.astype(np.float32, copy=False)
406
+
407
+ r = int(max(0, radius))
408
+ if r == 0:
409
+ c = _ssd_confidence(ref_m, cur_m, dx0, dy0, crop=crop)
410
+ return 0.0, 0.0, float(c)
411
+
412
+ # Apply current estimate once
413
+ cur0 = _shift_image(cur_m, float(dx0), float(dy0))
414
+
415
+ # Gradients once
416
+ rg = _grad_img(ref_m)
417
+ cg0 = _grad_img(cur0)
418
+
419
+ H, W = rg.shape[:2]
420
+ cfx = max(8, int(W * (1.0 - float(crop)) * 0.5))
421
+ cfy = max(8, int(H * (1.0 - float(crop)) * 0.5))
422
+ x0, x1 = cfx, W - cfx
423
+ y0, y1 = cfy, H - cfy
424
+
425
+ rgc = rg[y0:y1, x0:x1]
426
+ cgc0 = cg0[y0:y1, x0:x1]
427
+
428
+ # brute-force integer search
429
+ best = (0, 0)
430
+ best_ssd = float("inf")
431
+ ssds: dict[tuple[int, int], float] = {}
432
+
433
+ for j in range(-r, r + 1):
434
+ for i in range(-r, r + 1):
435
+ ssd = _ssd_confidence_prepared(rgc, cgc0, i, j)
436
+ ssds[(i, j)] = ssd
437
+ if ssd < best_ssd:
438
+ best_ssd = ssd
439
+ best = (i, j)
440
+
441
+ bx, by = best
442
+
443
+ # Subpixel quadratic fit (separable) if neighbors exist
444
+ def _quad_peak(vm, v0, vp):
445
+ denom = (vm - 2.0 * v0 + vp)
446
+ if abs(denom) < 1e-12:
447
+ return 0.0
448
+ return 0.5 * (vm - vp) / denom
449
+
450
+ dx_sub = 0.0
451
+ dy_sub = 0.0
452
+ if (bx - 1, by) in ssds and (bx + 1, by) in ssds:
453
+ dx_sub = _quad_peak(ssds[(bx - 1, by)], ssds[(bx, by)], ssds[(bx + 1, by)])
454
+ if (bx, by - 1) in ssds and (bx, by + 1) in ssds:
455
+ dy_sub = _quad_peak(ssds[(bx, by - 1)], ssds[(bx, by)], ssds[(bx, by + 1)])
456
+
457
+ dxr = float(bx + np.clip(dx_sub, -0.75, 0.75))
458
+ dyr = float(by + np.clip(dy_sub, -0.75, 0.75))
459
+
460
+ # Confidence: use your “sharpness” idea (median neighbor vs best)
461
+ neigh = [v for (k, v) in ssds.items() if k != (bx, by)]
462
+ neigh_med = float(np.median(np.asarray(neigh, np.float32))) if neigh else best_ssd
463
+ sharp = max(0.0, neigh_med - best_ssd)
464
+ conf = float(np.clip(sharp / max(1e-6, neigh_med), 0.0, 1.0))
465
+
466
+ return dxr, dyr, conf
467
+
468
+ def _bandpass(m: np.ndarray) -> np.ndarray:
469
+ """Illumination-robust image for tracking (float32)."""
470
+ m = m.astype(np.float32, copy=False)
471
+
472
+ # remove large-scale illumination (terminator gradient)
473
+ lo = cv2.GaussianBlur(m, (0, 0), 6.0)
474
+ hi = cv2.GaussianBlur(m, (0, 0), 1.2)
475
+ bp = hi - lo
476
+
477
+ # normalize
478
+ bp -= float(bp.mean())
479
+ s = float(bp.std()) + 1e-6
480
+ bp = bp / s
481
+
482
+ # window to reduce FFT edge artifacts
483
+ hann_y = np.hanning(bp.shape[0]).astype(np.float32)
484
+ hann_x = np.hanning(bp.shape[1]).astype(np.float32)
485
+ bp *= (hann_y[:, None] * hann_x[None, :])
486
+
487
+ return bp
488
+
489
+ def _reject_ap_outliers(ap_dx: np.ndarray, ap_dy: np.ndarray, ap_cf: np.ndarray, *, z: float = 3.5) -> np.ndarray:
490
+ """
491
+ Return a boolean mask of APs to keep based on MAD distance from median.
492
+ """
493
+ dx = np.asarray(ap_dx, np.float32)
494
+ dy = np.asarray(ap_dy, np.float32)
495
+ cf = np.asarray(ap_cf, np.float32)
496
+
497
+ good = cf > 0.15
498
+ if not np.any(good):
499
+ return good
500
+
501
+ dxg = dx[good]
502
+ dyg = dy[good]
503
+
504
+ mx = float(np.median(dxg))
505
+ my = float(np.median(dyg))
506
+
507
+ rx = np.abs(dxg - mx)
508
+ ry = np.abs(dyg - my)
509
+
510
+ madx = float(np.median(rx)) + 1e-6
511
+ mady = float(np.median(ry)) + 1e-6
512
+
513
+ zx = rx / madx
514
+ zy = ry / mady
515
+
516
+ keep_g = (zx < z) & (zy < z)
517
+ keep = np.zeros_like(good)
518
+ keep_idx = np.where(good)[0]
519
+ keep[keep_idx] = keep_g
520
+ return keep
521
+
522
+
523
+ def _coarse_surface_ref_locked(
524
+ source_obj,
525
+ *,
526
+ n: int,
527
+ roi,
528
+ roi_used=None,
529
+ debayer: bool,
530
+ to_rgb: bool,
531
+ bayer_pattern: Optional[str] = None,
532
+ progress_cb=None,
533
+ progress_every: int = 25,
534
+ # tuning:
535
+ down: int = 2,
536
+ template_size: int = 256,
537
+ search_radius: int = 96,
538
+ bandpass: bool = True,
539
+ # ✅ NEW: parallel coarse
540
+ workers: int | None = None,
541
+ stride: int = 16,
542
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
543
+ """
544
+ Surface coarse tracking that DOES NOT DRIFT:
545
+ - Locks to frame0 reference (in roi=roi_track coords).
546
+ - Uses NCC + subpixel phaseCorr.
547
+ - Optional parallelization by chunking time into segments of length=stride.
548
+ Each segment runs sequentially (keeps pred window), segments run in parallel.
549
+ """
550
+ if cv2 is None:
551
+ dx = np.zeros((n,), np.float32)
552
+ dy = np.zeros((n,), np.float32)
553
+ cc = np.ones((n,), np.float32)
554
+ return dx, dy, cc
555
+
556
+ dx = np.zeros((n,), dtype=np.float32)
557
+ dy = np.zeros((n,), dtype=np.float32)
558
+ cc = np.zeros((n,), dtype=np.float32)
559
+
560
+ def _downN(m: np.ndarray) -> np.ndarray:
561
+ if down <= 1:
562
+ return m.astype(np.float32, copy=False)
563
+ H, W = m.shape[:2]
564
+ return cv2.resize(
565
+ m,
566
+ (max(2, W // down), max(2, H // down)),
567
+ interpolation=cv2.INTER_AREA,
568
+ ).astype(np.float32, copy=False)
569
+
570
+ def _pick_anchor_center_ds(W: int, H: int) -> tuple[int, int]:
571
+ cx = W // 2
572
+ cy = H // 2
573
+ if roi_used is None or roi is None:
574
+ return int(cx), int(cy)
575
+ try:
576
+ xt, yt, wt, ht = [int(v) for v in roi]
577
+ xu, yu, wu, hu = [int(v) for v in roi_used]
578
+ cux = xu + (wu * 0.5)
579
+ cuy = yu + (hu * 0.5)
580
+ cx_full = cux - xt
581
+ cy_full = cuy - yt
582
+ cx = int(round(cx_full / max(1, int(down))))
583
+ cy = int(round(cy_full / max(1, int(down))))
584
+ cx = max(0, min(W - 1, cx))
585
+ cy = max(0, min(H - 1, cy))
586
+ except Exception:
587
+ pass
588
+ return int(cx), int(cy)
589
+
590
+ # ---------------------------
591
+ # Prep ref/template once
592
+ # ---------------------------
593
+ src0, owns0 = _ensure_source(source_obj, cache_items=2)
594
+ try:
595
+ img0 = _get_frame(src0, 0, roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern)
596
+
597
+ ref0 = _to_mono01(img0).astype(np.float32, copy=False)
598
+ ref0 = _downN(ref0)
599
+ ref0p = _bandpass(ref0) if bandpass else (ref0 - float(ref0.mean()))
600
+
601
+ H, W = ref0p.shape[:2]
602
+ ts = int(max(64, min(template_size, min(H, W) - 4)))
603
+ half = ts // 2
604
+
605
+ cx0, cy0 = _pick_anchor_center_ds(W, H)
606
+ rx0 = max(0, min(W - ts, cx0 - half))
607
+ ry0 = max(0, min(H - ts, cy0 - half))
608
+ ref_t = ref0p[ry0:ry0 + ts, rx0:rx0 + ts].copy()
609
+ finally:
610
+ if owns0:
611
+ try:
612
+ src0.close()
613
+ except Exception:
614
+ pass
615
+
616
+ dx[0] = 0.0
617
+ dy[0] = 0.0
618
+ cc[0] = 1.0
619
+
620
+ if progress_cb:
621
+ progress_cb(0, n, "Surface: coarse (ref-locked NCC+subpix)…")
622
+
623
+ # If no workers requested (or too small), fall back to sequential
624
+ if workers is None:
625
+ cpu = os.cpu_count() or 4
626
+ workers = max(1, min(cpu, 48))
627
+ workers = int(max(1, workers))
628
+ stride = int(max(4, stride))
629
+
630
+ # ---------------------------
631
+ # Core "one frame" matcher
632
+ # ---------------------------
633
+ def _match_one(curp: np.ndarray, pred_x: float, pred_y: float, r: int) -> tuple[float, float, float, float, float]:
634
+ # returns (mx_ds, my_ds, dx_full, dy_full, conf)
635
+ x0 = int(max(0, min(W - 1, pred_x - r)))
636
+ y0 = int(max(0, min(H - 1, pred_y - r)))
637
+ x1 = int(min(W, pred_x + r + ts))
638
+ y1 = int(min(H, pred_y + r + ts))
639
+
640
+ win = curp[y0:y1, x0:x1]
641
+ if win.shape[0] < ts or win.shape[1] < ts:
642
+ return float(pred_x), float(pred_y), 0.0, 0.0, 0.0
643
+
644
+ res = cv2.matchTemplate(win, ref_t, cv2.TM_CCOEFF_NORMED)
645
+ _, max_val, _, max_loc = cv2.minMaxLoc(res)
646
+ conf_ncc = float(np.clip(max_val, 0.0, 1.0))
647
+
648
+ mx_ds = float(x0 + max_loc[0])
649
+ my_ds = float(y0 + max_loc[1])
650
+
651
+ # subpix refine on the matched patch
652
+ mx_i = int(round(mx_ds))
653
+ my_i = int(round(my_ds))
654
+ cur_t = curp[my_i:my_i + ts, mx_i:mx_i + ts]
655
+ if cur_t.shape == ref_t.shape:
656
+ (sdx, sdy), resp = cv2.phaseCorrelate(ref_t.astype(np.float32), cur_t.astype(np.float32))
657
+ sub_dx = float(sdx)
658
+ sub_dy = float(sdy)
659
+ conf_pc = float(np.clip(resp, 0.0, 1.0))
660
+ else:
661
+ sub_dx = 0.0
662
+ sub_dy = 0.0
663
+ conf_pc = 0.0
664
+
665
+ dx_ds = float(rx0 - mx_ds) + sub_dx
666
+ dy_ds = float(ry0 - my_ds) + sub_dy
667
+ dx_full = float(dx_ds * down)
668
+ dy_full = float(dy_ds * down)
669
+
670
+ conf = float(np.clip(0.65 * conf_ncc + 0.35 * conf_pc, 0.0, 1.0))
671
+ return float(mx_ds), float(my_ds), dx_full, dy_full, conf
672
+
673
+ # ---------------------------
674
+ # Keyframe boundary pass (sequential)
675
+ # ---------------------------
676
+ boundaries = list(range(0, n, stride))
677
+ start_pred = {} # b -> (pred_x, pred_y)
678
+ start_pred[0] = (float(rx0), float(ry0))
679
+
680
+ # We use a slightly larger radius for boundary frames to be extra safe
681
+ r_key = int(max(16, int(search_radius) * 2))
682
+
683
+ srck, ownsk = _ensure_source(source_obj, cache_items=2)
684
+ try:
685
+ pred_x, pred_y = float(rx0), float(ry0)
686
+ for b in boundaries[1:]:
687
+ img = _get_frame(srck, b, roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern)
688
+
689
+ cur = _to_mono01(img).astype(np.float32, copy=False)
690
+ cur = _downN(cur)
691
+ curp = _bandpass(cur) if bandpass else (cur - float(cur.mean()))
692
+
693
+ mx_ds, my_ds, dx_b, dy_b, conf_b = _match_one(curp, pred_x, pred_y, r_key)
694
+
695
+ # store boundary predictor (template top-left in this frame)
696
+ start_pred[b] = (mx_ds, my_ds)
697
+
698
+ # update for next boundary
699
+ pred_x, pred_y = mx_ds, my_ds
700
+
701
+ # also fill boundary output immediately (optional but nice)
702
+ dx[b] = dx_b
703
+ dy[b] = dy_b
704
+ cc[b] = conf_b
705
+ if conf_b < 0.15 and b > 0:
706
+ dx[b] = dx[b - 1]
707
+ dy[b] = dy[b - 1]
708
+ finally:
709
+ if ownsk:
710
+ try:
711
+ srck.close()
712
+ except Exception:
713
+ pass
714
+
715
+ # ---------------------------
716
+ # Parallel per-chunk scan (each chunk sequential)
717
+ # ---------------------------
718
+ r = int(max(16, search_radius))
719
+
720
+ def _run_chunk(b: int, e: int) -> int:
721
+ src, owns = _ensure_source(source_obj, cache_items=0)
722
+ try:
723
+ pred_x, pred_y = start_pred.get(b, (float(rx0), float(ry0)))
724
+ # if boundary already computed above, keep it; start after b
725
+ i0 = b
726
+ if b in start_pred and b != 0:
727
+ i0 = b + 1 # boundary already solved with r_key
728
+
729
+ if i0 == 0:
730
+ i0 = 1
731
+ for i in range(i0, e):
732
+ if i in start_pred:
733
+ pred_x, pred_y = start_pred[i]
734
+ continue
735
+
736
+ img = _get_frame(src, i, roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern)
737
+ cur = _to_mono01(img).astype(np.float32, copy=False)
738
+ cur = _downN(cur)
739
+ curp = _bandpass(cur) if bandpass else (cur - float(cur.mean()))
740
+
741
+ mx_ds, my_ds, dx_i, dy_i, conf_i = _match_one(curp, pred_x, pred_y, r)
742
+
743
+ dx[i] = dx_i
744
+ dy[i] = dy_i
745
+ cc[i] = conf_i
746
+
747
+ pred_x, pred_y = mx_ds, my_ds
748
+
749
+ if conf_i < 0.15 and i > 0:
750
+ dx[i] = dx[i - 1]
751
+ dy[i] = dy[i - 1]
752
+ return (e - b)
753
+ finally:
754
+ if owns:
755
+ try:
756
+ src.close()
757
+ except Exception:
758
+ pass
759
+
760
+ if workers <= 1 or n <= stride * 2:
761
+ # small job: just do sequential scan exactly like before
762
+ src, owns = _ensure_source(source_obj, cache_items=2)
763
+ try:
764
+ pred_x, pred_y = float(rx0), float(ry0)
765
+ for i in range(1, n):
766
+ img = _get_frame(src, i, roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern)
767
+ cur = _to_mono01(img).astype(np.float32, copy=False)
768
+ cur = _downN(cur)
769
+ curp = _bandpass(cur) if bandpass else (cur - float(cur.mean()))
770
+
771
+ mx_ds, my_ds, dx_i, dy_i, conf_i = _match_one(curp, pred_x, pred_y, r)
772
+ dx[i] = dx_i
773
+ dy[i] = dy_i
774
+ cc[i] = conf_i
775
+ pred_x, pred_y = mx_ds, my_ds
776
+
777
+ if conf_i < 0.15:
778
+ dx[i] = dx[i - 1]
779
+ dy[i] = dy[i - 1]
780
+
781
+ if progress_cb and (i % int(max(1, progress_every)) == 0 or i == n - 1):
782
+ progress_cb(i, n, "Surface: coarse (ref-locked NCC+subpix)…")
783
+ finally:
784
+ if owns:
785
+ try:
786
+ src.close()
787
+ except Exception:
788
+ pass
789
+ return dx, dy, cc
790
+
791
+ # Parallel chunks
792
+ done = 0
793
+ with ThreadPoolExecutor(max_workers=workers) as ex:
794
+ futs = []
795
+ for b in boundaries:
796
+ e = min(n, b + stride)
797
+ futs.append(ex.submit(_run_chunk, b, e))
798
+
799
+ for fut in as_completed(futs):
800
+ done += int(fut.result())
801
+ if progress_cb:
802
+ # best-effort: done is "frames processed" not exact index
803
+ progress_cb(min(done, n - 1), n, "Surface: coarse (ref-locked NCC+subpix)…")
804
+
805
+ return dx, dy, cc
806
+
807
+
808
+ def _shift_image(img01: np.ndarray, dx: float, dy: float) -> np.ndarray:
809
+ """
810
+ Shift image by (dx,dy) in pixel units. Positive dx shifts right, positive dy shifts down.
811
+ Uses cv2.warpAffine if available; else nearest-ish roll (wrap) fallback.
812
+ """
813
+ if abs(dx) < 1e-6 and abs(dy) < 1e-6:
814
+ return img01
815
+
816
+ if cv2 is not None:
817
+ # border replicate is usually better than constant black for planetary
818
+ h, w = img01.shape[:2]
819
+ M = np.array([[1.0, 0.0, dx],
820
+ [0.0, 1.0, dy]], dtype=np.float32)
821
+ if img01.ndim == 2:
822
+ return cv2.warpAffine(img01, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
823
+ else:
824
+ return cv2.warpAffine(img01, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
825
+ # very rough fallback (wraps!)
826
+ rx = int(round(dx))
827
+ ry = int(round(dy))
828
+ out = np.roll(img01, shift=ry, axis=0)
829
+ out = np.roll(out, shift=rx, axis=1)
830
+ return out
831
+
832
+ def _downsample_mono01(img01: np.ndarray, max_dim: int = 512) -> np.ndarray:
833
+ """
834
+ Convert to mono and downsample for analysis/tracking. Returns float32 in [0,1].
835
+ """
836
+ m = _to_mono01(img01).astype(np.float32, copy=False)
837
+ H, W = m.shape[:2]
838
+ mx = int(max(1, max_dim))
839
+ if max(H, W) <= mx:
840
+ return m
841
+
842
+ if cv2 is None:
843
+ # crude fallback
844
+ scale = mx / float(max(H, W))
845
+ nh = max(2, int(round(H * scale)))
846
+ nw = max(2, int(round(W * scale)))
847
+ # nearest-ish
848
+ ys = (np.linspace(0, H - 1, nh)).astype(np.int32)
849
+ xs = (np.linspace(0, W - 1, nw)).astype(np.int32)
850
+ return m[ys[:, None], xs[None, :]].astype(np.float32)
851
+
852
+ scale = mx / float(max(H, W))
853
+ nh = max(2, int(round(H * scale)))
854
+ nw = max(2, int(round(W * scale)))
855
+ return cv2.resize(m, (nw, nh), interpolation=cv2.INTER_AREA).astype(np.float32, copy=False)
856
+
857
+
858
+ def _phase_corr_shift(ref_m: np.ndarray, cur_m: np.ndarray) -> tuple[float, float, float]:
859
+ """
860
+ Returns (dx, dy, response) such that shifting cur by (dx,dy) aligns to ref.
861
+ Uses cv2.phaseCorrelate if available.
862
+ """
863
+ if cv2 is None:
864
+ return 0.0, 0.0, 1.0
865
+
866
+ # phaseCorrelate expects float32/float64
867
+ ref = ref_m.astype(np.float32, copy=False)
868
+ cur = cur_m.astype(np.float32, copy=False)
869
+ (dx, dy), resp = cv2.phaseCorrelate(ref, cur) # shift cur -> ref
870
+ return float(dx), float(dy), float(resp)
871
+
872
+ def _ensure_source(source, cache_items: int = 10) -> tuple[PlanetaryFrameSource, bool]:
873
+ """
874
+ Returns (src, owns_src)
875
+
876
+ Accepts:
877
+ - PlanetaryFrameSource-like object (duck typed: get_frame/meta/close)
878
+ - path string
879
+ - list/tuple of paths
880
+ """
881
+ # Already an opened source-like object
882
+ if source is not None and hasattr(source, "get_frame") and hasattr(source, "meta") and hasattr(source, "close"):
883
+ return source, False
884
+
885
+ # allow tuple -> list
886
+ if isinstance(source, tuple):
887
+ source = list(source)
888
+
889
+ src = open_planetary_source(source, cache_items=cache_items)
890
+ return src, True
891
+
892
+ def stack_ser(
893
+ source: str | list[str] | PlanetaryFrameSource,
894
+ *,
895
+ roi=None,
896
+ debayer: bool = True,
897
+ keep_percent: float = 20.0,
898
+ track_mode: str = "planetary",
899
+ surface_anchor=None,
900
+ to_rgb: bool = False, # ✅ add this
901
+ bayer_pattern: Optional[str] = None, # ✅ strongly recommended since dialog passes it
902
+ analysis: AnalyzeResult | None = None,
903
+ local_warp: bool = True,
904
+ max_dim: int = 512,
905
+ progress_cb=None,
906
+ cache_items: int = 10,
907
+ workers: int | None = None,
908
+ chunk_size: int | None = None,
909
+ # ✅ NEW drizzle knobs
910
+ drizzle_scale: float = 1.0,
911
+ drizzle_pixfrac: float = 0.80,
912
+ drizzle_kernel: str = "gaussian",
913
+ drizzle_sigma: float = 0.0,
914
+
915
+ ) -> tuple[np.ndarray, dict]:
916
+ source_obj = source
917
+
918
+ # ---- Worker count ----
919
+ if workers is None:
920
+ cpu = os.cpu_count() or 4
921
+ workers = max(1, min(cpu, 48))
922
+
923
+ if cv2 is not None:
924
+ try:
925
+ cv2.setNumThreads(1)
926
+ except Exception:
927
+ pass
928
+
929
+ drizzle_scale = float(drizzle_scale)
930
+ drizzle_on = drizzle_scale > 1.0001
931
+ drizzle_pixfrac = float(drizzle_pixfrac)
932
+ drizzle_kernel = str(drizzle_kernel).strip().lower()
933
+ if drizzle_kernel not in ("square", "circle", "gaussian"):
934
+ drizzle_kernel = "gaussian"
935
+ drizzle_sigma = float(drizzle_sigma)
936
+
937
+ # ---- Open once to get meta + first frame shape ----
938
+ src0, owns0 = _ensure_source(source_obj, cache_items=cache_items)
939
+ try:
940
+ n = int(src0.meta.frames)
941
+ keep_percent = max(0.1, min(100.0, float(keep_percent)))
942
+ k = max(1, int(round(n * (keep_percent / 100.0))))
943
+
944
+ if analysis is None or analysis.ref_image is None or analysis.ap_centers is None:
945
+ raise ValueError("stack_ser expects analysis with ref_image + ap_centers (run Analyze first).")
946
+
947
+ order = np.asarray(analysis.order, np.int32)
948
+ keep_idx = order[:k].astype(np.int32, copy=False)
949
+
950
+ # reference / APs
951
+ ref_img = analysis.ref_image.astype(np.float32, copy=False)
952
+ ref_m = _to_mono01(ref_img).astype(np.float32, copy=False)
953
+ ap_centers_all = np.asarray(analysis.ap_centers, np.int32)
954
+ ap_size = int(getattr(analysis, "ap_size", 64) or 64)
955
+
956
+ # frame shape for accumulator
957
+ first = _get_frame(src0, int(keep_idx[0]), roi=roi, debayer=debayer, to_float01=True, force_rgb=False, bayer_pattern=bayer_pattern)
958
+ acc_shape = first.shape # (H,W) or (H,W,3)
959
+ finally:
960
+ if owns0:
961
+ try:
962
+ src0.close()
963
+ except Exception:
964
+ pass
965
+
966
+ # ---- Progress aggregation (thread-safe) ----
967
+ done_lock = threading.Lock()
968
+ done_ct = 0
969
+ total_ct = int(len(keep_idx))
970
+
971
+ def _bump_progress(delta: int, phase: str = "Stack"):
972
+ nonlocal done_ct
973
+ if progress_cb is None:
974
+ return
975
+ with done_lock:
976
+ done_ct += int(delta)
977
+ d = done_ct
978
+ progress_cb(d, total_ct, phase)
979
+
980
+ # ---- Chunking ----
981
+ idx_list = keep_idx.tolist()
982
+ if chunk_size is None:
983
+ chunk_size = max(8, int(np.ceil(len(idx_list) / float(workers * 2))))
984
+ chunks: list[list[int]] = [idx_list[i:i + chunk_size] for i in range(0, len(idx_list), chunk_size)]
985
+
986
+ if progress_cb:
987
+ progress_cb(0, total_ct, "Stack")
988
+
989
+ # ---- drizzle helpers ----
990
+ if drizzle_on:
991
+ from setiastro.saspro.legacy.numba_utils import (
992
+ drizzle_deposit_numba_kernel_mono,
993
+ drizzle_deposit_color_kernel,
994
+ finalize_drizzle_2d,
995
+ finalize_drizzle_3d,
996
+ )
997
+
998
+ # map kernel string -> code used by your numba
999
+ kernel_code = {"square": 0, "circle": 1, "gaussian": 2}[drizzle_kernel]
1000
+
1001
+ # If gaussian sigma isn't provided, use something tied to pixfrac.
1002
+ # Your numba interprets gaussian sigma as "sigma_out", and also enforces >= drop_shrink*0.5.
1003
+ if drizzle_sigma <= 1e-9:
1004
+ # a good practical default: sigma ~ pixfrac*0.5
1005
+ drizzle_sigma_eff = max(1e-3, float(drizzle_pixfrac) * 0.5)
1006
+ else:
1007
+ drizzle_sigma_eff = drizzle_sigma
1008
+
1009
+ H, W = int(acc_shape[0]), int(acc_shape[1])
1010
+ outH = int(round(H * drizzle_scale))
1011
+ outW = int(round(W * drizzle_scale))
1012
+
1013
+ # Identity transform from input pixels -> aligned/reference pixel coords
1014
+ # drizzle_factor applies the scale.
1015
+ T = np.zeros((2, 3), dtype=np.float32)
1016
+ T[0, 0] = 1.0
1017
+ T[1, 1] = 1.0
1018
+
1019
+ # ---- Worker: accumulate its own sum OR its own drizzle buffers ----
1020
+ def _stack_chunk(chunk: list[int]):
1021
+ src, owns = _ensure_source(source_obj, cache_items=0)
1022
+ try:
1023
+ if drizzle_on:
1024
+ if len(acc_shape) == 2:
1025
+ dbuf = np.zeros((outH, outW), dtype=np.float32)
1026
+ cbuf = np.zeros((outH, outW), dtype=np.float32)
1027
+ else:
1028
+ dbuf = np.zeros((outH, outW, acc_shape[2]), dtype=np.float32)
1029
+ cbuf = np.zeros((outH, outW, acc_shape[2]), dtype=np.float32)
1030
+ else:
1031
+ acc = np.zeros(acc_shape, dtype=np.float32)
1032
+ wacc = 0.0
1033
+
1034
+ for i in chunk:
1035
+ img = _get_frame(src, int(i), roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern).astype(np.float32, copy=False)
1036
+
1037
+ # Global prior (from Analyze)
1038
+ gdx = float(analysis.dx[int(i)]) if (analysis.dx is not None) else 0.0
1039
+ gdy = float(analysis.dy[int(i)]) if (analysis.dy is not None) else 0.0
1040
+
1041
+ # Global prior always first
1042
+ warped_g = _shift_image(img, gdx, gdy)
1043
+
1044
+ if cv2 is None or (not local_warp):
1045
+ warped = warped_g
1046
+ else:
1047
+ cur_m_g = _to_mono01(warped_g).astype(np.float32, copy=False)
1048
+
1049
+ ap_rdx, ap_rdy, ap_resp = _ap_phase_shifts_per_ap(
1050
+ ref_m, cur_m_g,
1051
+ ap_centers=ap_centers_all,
1052
+ ap_size=ap_size,
1053
+ max_dim=max_dim,
1054
+ )
1055
+ ap_cf = np.clip(ap_resp.astype(np.float32, copy=False), 0.0, 1.0)
1056
+
1057
+ keep = _reject_ap_outliers(ap_rdx, ap_rdy, ap_cf, z=3.5)
1058
+ if np.any(keep):
1059
+ ap_centers = ap_centers_all[keep]
1060
+ ap_dx_k = ap_rdx[keep]
1061
+ ap_dy_k = ap_rdy[keep]
1062
+ ap_cf_k = ap_cf[keep]
1063
+
1064
+ dx_field, dy_field = _dense_field_from_ap_shifts(
1065
+ warped_g.shape[0], warped_g.shape[1],
1066
+ ap_centers, ap_dx_k, ap_dy_k, ap_cf_k,
1067
+ grid=32, power=2.0, conf_floor=0.15,
1068
+ radius=float(ap_size) * 3.0,
1069
+ )
1070
+ warped = _warp_by_dense_field(warped_g, dx_field, dy_field)
1071
+ else:
1072
+ warped = warped_g
1073
+
1074
+ if drizzle_on:
1075
+ # deposit aligned frame into drizzle buffers
1076
+ fw = 1.0 # frame_weight (could later use quality weights)
1077
+ if warped.ndim == 2:
1078
+ drizzle_deposit_numba_kernel_mono(
1079
+ warped, T, dbuf, cbuf,
1080
+ drizzle_factor=drizzle_scale,
1081
+ drop_shrink=drizzle_pixfrac,
1082
+ frame_weight=fw,
1083
+ kernel_code=kernel_code,
1084
+ gaussian_sigma_or_radius=drizzle_sigma_eff,
1085
+ )
1086
+ else:
1087
+ drizzle_deposit_color_kernel(
1088
+ warped, T, dbuf, cbuf,
1089
+ drizzle_factor=drizzle_scale,
1090
+ drop_shrink=drizzle_pixfrac,
1091
+ frame_weight=fw,
1092
+ kernel_code=kernel_code,
1093
+ gaussian_sigma_or_radius=drizzle_sigma_eff,
1094
+ )
1095
+ else:
1096
+ acc += warped
1097
+ wacc += 1.0
1098
+
1099
+ _bump_progress(len(chunk), "Stack")
1100
+
1101
+ if drizzle_on:
1102
+ return dbuf, cbuf
1103
+ return acc, wacc
1104
+
1105
+ finally:
1106
+ if owns:
1107
+ try:
1108
+ src.close()
1109
+ except Exception:
1110
+ pass
1111
+
1112
+ # ---- Parallel run + reduce ----
1113
+ if drizzle_on:
1114
+ # reduce drizzle buffers
1115
+ if len(acc_shape) == 2:
1116
+ dbuf_total = np.zeros((outH, outW), dtype=np.float32)
1117
+ cbuf_total = np.zeros((outH, outW), dtype=np.float32)
1118
+ else:
1119
+ dbuf_total = np.zeros((outH, outW, acc_shape[2]), dtype=np.float32)
1120
+ cbuf_total = np.zeros((outH, outW, acc_shape[2]), dtype=np.float32)
1121
+
1122
+ with ThreadPoolExecutor(max_workers=workers) as ex:
1123
+ futs = [ex.submit(_stack_chunk, c) for c in chunks if c]
1124
+ for fut in as_completed(futs):
1125
+ db, cb = fut.result()
1126
+ dbuf_total += db
1127
+ cbuf_total += cb
1128
+
1129
+ # finalize
1130
+ if len(acc_shape) == 2:
1131
+ out = np.zeros((outH, outW), dtype=np.float32)
1132
+ finalize_drizzle_2d(dbuf_total, cbuf_total, out)
1133
+ else:
1134
+ out = np.zeros((outH, outW, acc_shape[2]), dtype=np.float32)
1135
+ finalize_drizzle_3d(dbuf_total, cbuf_total, out)
1136
+
1137
+ out = np.clip(out, 0.0, 1.0).astype(np.float32, copy=False)
1138
+
1139
+ else:
1140
+ acc_total = np.zeros(acc_shape, dtype=np.float32)
1141
+ wacc_total = 0.0
1142
+
1143
+ with ThreadPoolExecutor(max_workers=workers) as ex:
1144
+ futs = [ex.submit(_stack_chunk, c) for c in chunks if c]
1145
+ for fut in as_completed(futs):
1146
+ acc_c, w_c = fut.result()
1147
+ acc_total += acc_c
1148
+ wacc_total += float(w_c)
1149
+
1150
+ out = np.clip(acc_total / max(1e-6, wacc_total), 0.0, 1.0).astype(np.float32, copy=False)
1151
+
1152
+ diag = {
1153
+ "frames_total": int(n),
1154
+ "frames_kept": int(len(keep_idx)),
1155
+ "roi_used": roi,
1156
+ "track_mode": track_mode,
1157
+ "local_warp": bool(local_warp),
1158
+ "workers": int(workers),
1159
+ "chunk_size": int(chunk_size),
1160
+ "drizzle_scale": float(drizzle_scale),
1161
+ "drizzle_pixfrac": float(drizzle_pixfrac),
1162
+ "drizzle_kernel": str(drizzle_kernel),
1163
+ "drizzle_sigma": float(drizzle_sigma),
1164
+ }
1165
+ return out, diag
1166
+
1167
+ def _build_reference(
1168
+ src: PlanetaryFrameSource,
1169
+ *,
1170
+ order: np.ndarray,
1171
+ roi,
1172
+ debayer: bool,
1173
+ to_rgb: bool,
1174
+ ref_mode: str,
1175
+ ref_count: int,
1176
+ bayer_pattern=None,
1177
+ ) -> np.ndarray:
1178
+ """
1179
+ ref_mode:
1180
+ - "best_frame": return best single frame
1181
+ - "best_stack": return mean of best ref_count frames
1182
+ """
1183
+ best_idx = int(order[0])
1184
+ f0 = _get_frame(src, best_idx, roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern)
1185
+ if ref_mode != "best_stack" or ref_count <= 1:
1186
+ return f0.astype(np.float32, copy=False)
1187
+
1188
+ k = int(max(2, min(ref_count, len(order))))
1189
+ acc = np.zeros_like(f0, dtype=np.float32)
1190
+ for j in range(k):
1191
+ idx = int(order[j])
1192
+ fr = _get_frame(src, idx, roi=roi, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb), bayer_pattern=bayer_pattern)
1193
+ acc += fr.astype(np.float32, copy=False)
1194
+ ref = acc / float(k)
1195
+ return np.clip(ref, 0.0, 1.0).astype(np.float32, copy=False)
1196
+
1197
+ def _cfg_get_source(cfg) -> Any:
1198
+ """
1199
+ Back-compat: prefer cfg.source (new), else cfg.ser_path (old).
1200
+ cfg.source may be:
1201
+ - path string (ser/avi/mp4/etc)
1202
+ - list of image paths
1203
+ - PlanetaryFrameSource
1204
+ """
1205
+ src = getattr(cfg, "source", None)
1206
+ if src is not None and src != "":
1207
+ return src
1208
+ return getattr(cfg, "ser_path", None)
1209
+
1210
+ def analyze_ser(
1211
+ cfg: SERStackConfig,
1212
+ *,
1213
+ debayer: bool = True,
1214
+ to_rgb: bool = False,
1215
+ smooth_sigma: float = 1.5, # kept for API compat
1216
+ thresh_pct: float = 92.0, # kept for API compat
1217
+ ref_mode: str = "best_frame", # "best_frame" or "best_stack"
1218
+ bayer_pattern: Optional[str] = None,
1219
+ ref_count: int = 5,
1220
+ max_dim: int = 512,
1221
+ progress_cb=None,
1222
+ workers: Optional[int] = None,
1223
+ ) -> AnalyzeResult:
1224
+ """
1225
+ Parallel analyze for *any* PlanetaryFrameSource (SER/AVI/MP4/images/sequence).
1226
+ - Pass 1: quality for every frame
1227
+ - Build reference:
1228
+ - planetary: best frame or best-N stack
1229
+ - surface: frame 0 (chronological anchor)
1230
+ - Autoplace APs (always)
1231
+ - Pass 2:
1232
+ - planetary: AP-based shift directly
1233
+ - surface:
1234
+ (A) coarse drift stabilization via ref-locked NCC+subpix (on a larger tracking ROI),
1235
+ (B) AP search+refine that follows coarse, with outlier rejection,
1236
+ (C) robust median -> final dx/dy/conf
1237
+ """
1238
+
1239
+ source_obj = _cfg_get_source(cfg)
1240
+ bpat = bayer_pattern or _cfg_bayer_pattern(cfg)
1241
+
1242
+ if not source_obj:
1243
+ raise ValueError("SERStackConfig.source/ser_path is empty")
1244
+
1245
+ # ---- open source + meta (single open) ----
1246
+ src0, owns0 = _ensure_source(source_obj, cache_items=2)
1247
+ try:
1248
+ meta = src0.meta
1249
+ base_roi = cfg.roi
1250
+ if base_roi is not None:
1251
+ base_roi = _clamp_roi_in_bounds(base_roi, meta.width, meta.height)
1252
+ n = int(meta.frames)
1253
+ if n <= 0:
1254
+ raise ValueError("Source contains no frames")
1255
+ src_w = int(meta.width)
1256
+ src_h = int(meta.height)
1257
+ finally:
1258
+ if owns0:
1259
+ try:
1260
+ src0.close()
1261
+ except Exception:
1262
+ pass
1263
+
1264
+ # ---- Worker count ----
1265
+ if workers is None:
1266
+ cpu = os.cpu_count() or 4
1267
+ workers = max(1, min(cpu, 48))
1268
+
1269
+ if cv2 is not None:
1270
+ try:
1271
+ cv2.setNumThreads(1)
1272
+ except Exception:
1273
+ pass
1274
+
1275
+ # ---- Surface tracking ROI (IMPORTANT for big drift) ----
1276
+ def _surface_tracking_roi() -> Optional[Tuple[int, int, int, int]]:
1277
+ if base_roi is None:
1278
+ return None # full frame
1279
+ margin = int(getattr(cfg, "surface_track_margin", 256))
1280
+ x, y, w, h = [int(v) for v in base_roi]
1281
+ x0 = max(0, x - margin)
1282
+ y0 = max(0, y - margin)
1283
+ x1 = min(src_w, x + w + margin)
1284
+ y1 = min(src_h, y + h + margin)
1285
+ return _clamp_roi_in_bounds((x0, y0, x1 - x0, y1 - y0), src_w, src_h)
1286
+
1287
+ roi_track = _surface_tracking_roi() if cfg.track_mode == "surface" else base_roi
1288
+ roi_used = base_roi # APs and final ref are in this coordinate system
1289
+
1290
+ # -------------------------------------------------------------------------
1291
+ # Pass 1: quality (use roi_used)
1292
+ # -------------------------------------------------------------------------
1293
+ quality = np.zeros((n,), dtype=np.float32)
1294
+ idxs = np.arange(n, dtype=np.int32)
1295
+ n_chunks = max(5, int(workers) * int(getattr(cfg, "progress_chunk_factor", 5)))
1296
+ n_chunks = max(1, min(int(n), n_chunks))
1297
+ chunks = np.array_split(idxs, n_chunks)
1298
+
1299
+ if progress_cb:
1300
+ progress_cb(0, n, "Quality")
1301
+
1302
+ def _q_chunk(chunk: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
1303
+ out_i: list[int] = []
1304
+ out_q: list[float] = []
1305
+ src, owns = _ensure_source(source_obj, cache_items=0)
1306
+ try:
1307
+ for i in chunk.tolist():
1308
+ img = _get_frame(
1309
+ src, int(i),
1310
+ roi=roi_used,
1311
+ debayer=debayer,
1312
+ to_float01=True,
1313
+ force_rgb=bool(to_rgb),
1314
+ bayer_pattern=bpat,
1315
+ )
1316
+ m = _downsample_mono01(img, max_dim=max_dim)
1317
+
1318
+ if cv2 is not None:
1319
+ lap = cv2.Laplacian(m, cv2.CV_32F, ksize=3)
1320
+ q = float(np.mean(np.abs(lap)))
1321
+ else:
1322
+ q = float(
1323
+ np.abs(m[:, 1:] - m[:, :-1]).mean() +
1324
+ np.abs(m[1:, :] - m[:-1, :]).mean()
1325
+ )
1326
+ out_i.append(int(i))
1327
+ out_q.append(q)
1328
+ finally:
1329
+ if owns:
1330
+ try:
1331
+ src.close()
1332
+ except Exception:
1333
+ pass
1334
+ return np.asarray(out_i, np.int32), np.asarray(out_q, np.float32)
1335
+
1336
+ done_ct = 0
1337
+ with ThreadPoolExecutor(max_workers=workers) as ex:
1338
+ futs = [ex.submit(_q_chunk, c) for c in chunks if c.size > 0]
1339
+ for fut in as_completed(futs):
1340
+ ii, qq = fut.result()
1341
+ quality[ii] = qq
1342
+ done_ct += int(ii.size)
1343
+ if progress_cb:
1344
+ progress_cb(done_ct, n, "Quality")
1345
+
1346
+ order = np.argsort(-quality).astype(np.int32, copy=False)
1347
+
1348
+ # -------------------------------------------------------------------------
1349
+ # Build reference
1350
+ # -------------------------------------------------------------------------
1351
+ ref_count = int(max(1, min(int(ref_count), n)))
1352
+ ref_mode = "best_stack" if ref_mode == "best_stack" else "best_frame"
1353
+
1354
+ src_ref, owns_ref = _ensure_source(source_obj, cache_items=2)
1355
+ if progress_cb:
1356
+ progress_cb(0, n, f"Building reference ({ref_mode}, N={ref_count})…")
1357
+ try:
1358
+ if cfg.track_mode == "surface":
1359
+ # Surface ref must be frame 0 in roi_used coords
1360
+ ref_img = _get_frame(
1361
+ src_ref, 0,
1362
+ roi=roi_used, debayer=debayer, to_float01=True, force_rgb=bool(to_rgb),
1363
+ bayer_pattern=bpat,
1364
+ ).astype(np.float32, copy=False)
1365
+
1366
+ ref_mode = "first_frame"
1367
+ ref_count = 1
1368
+ else:
1369
+ ref_img = _build_reference(
1370
+ src_ref,
1371
+ order=order,
1372
+ roi=roi_used,
1373
+ debayer=debayer,
1374
+ to_rgb=to_rgb,
1375
+ ref_mode=ref_mode,
1376
+ ref_count=ref_count,
1377
+ bayer_pattern=bpat, # ✅ add this
1378
+ ).astype(np.float32, copy=False)
1379
+
1380
+ finally:
1381
+ if owns_ref:
1382
+ try:
1383
+ src_ref.close()
1384
+ except Exception:
1385
+ pass
1386
+
1387
+ # -------------------------------------------------------------------------
1388
+ # Autoplace APs (always)
1389
+ # -------------------------------------------------------------------------
1390
+ if progress_cb:
1391
+ progress_cb(0, n, "Placing alignment points…")
1392
+
1393
+ ap_size = int(getattr(cfg, "ap_size", 64) or 64)
1394
+ ap_centers = _autoplace_aps(
1395
+ ref_img,
1396
+ ap_size=ap_size,
1397
+ ap_spacing=int(getattr(cfg, "ap_spacing", 48)),
1398
+ ap_min_mean=float(getattr(cfg, "ap_min_mean", 0.03)),
1399
+ )
1400
+
1401
+ # -------------------------------------------------------------------------
1402
+ # Pass 2: shifts/conf
1403
+ # -------------------------------------------------------------------------
1404
+ dx = np.zeros((n,), dtype=np.float32)
1405
+ dy = np.zeros((n,), dtype=np.float32)
1406
+ conf = np.ones((n,), dtype=np.float32)
1407
+ coarse_conf: Optional[np.ndarray] = None
1408
+
1409
+ if cfg.track_mode == "off" or cv2 is None:
1410
+ return AnalyzeResult(
1411
+ frames_total=n,
1412
+ roi_used=roi_used,
1413
+ track_mode=cfg.track_mode,
1414
+ quality=quality,
1415
+ dx=dx,
1416
+ dy=dy,
1417
+ conf=conf,
1418
+ order=order,
1419
+ ref_mode=ref_mode,
1420
+ ref_count=ref_count,
1421
+ ref_image=ref_img,
1422
+ ap_centers=ap_centers,
1423
+ ap_size=ap_size,
1424
+ ap_multiscale=bool(getattr(cfg, "ap_multiscale", False)),
1425
+ coarse_conf=None,
1426
+ )
1427
+
1428
+ ref_m_full = _to_mono01(ref_img).astype(np.float32, copy=False)
1429
+ use_multiscale = bool(getattr(cfg, "ap_multiscale", False))
1430
+
1431
+ # ---- surface coarse drift (ref-locked) ----
1432
+ if cfg.track_mode == "surface":
1433
+ coarse_conf = np.zeros((n,), dtype=np.float32)
1434
+ if progress_cb:
1435
+ progress_cb(0, n, "Surface: coarse drift (ref-locked NCC+subpix)…")
1436
+
1437
+ dx_chain, dy_chain, cc_chain = _coarse_surface_ref_locked(
1438
+ source_obj,
1439
+ n=n,
1440
+ roi=roi_track,
1441
+ roi_used=roi_used, # ✅ NEW
1442
+ debayer=debayer,
1443
+ to_rgb=to_rgb,
1444
+ bayer_pattern=bpat,
1445
+ progress_cb=progress_cb,
1446
+ progress_every=25,
1447
+ down=2,
1448
+ template_size=256,
1449
+ search_radius=96,
1450
+ bandpass=True,
1451
+ workers=min(workers, 8), # coarse doesn’t need 48; 4–8 is usually ideal
1452
+ stride=16, # 8–32 typical
1453
+ )
1454
+ dx[:] = dx_chain
1455
+ dy[:] = dy_chain
1456
+ coarse_conf[:] = cc_chain
1457
+
1458
+ # ---- chunked refine ----
1459
+ idxs2 = np.arange(n, dtype=np.int32)
1460
+
1461
+ # More/smaller chunks => progress updates sooner (futures complete more frequently)
1462
+ chunk_factor = int(getattr(cfg, "progress_chunk_factor", 5)) # optional knob
1463
+ min_chunks = 5
1464
+ n_chunks2 = max(min_chunks, int(workers) * chunk_factor)
1465
+ n_chunks2 = max(1, min(int(n), n_chunks2))
1466
+
1467
+ chunks2 = np.array_split(idxs2, n_chunks2)
1468
+
1469
+ if progress_cb:
1470
+ progress_cb(0, n, "SSD Refine")
1471
+
1472
+ if cfg.track_mode == "surface":
1473
+ # FAST surface refine:
1474
+ # - use coarse dx/dy from ref-locked tracker
1475
+ # - apply coarse shift to current mono frame
1476
+ # - compute residual per-AP phase shifts (NO SEARCH)
1477
+ # - final dx/dy = coarse + median(residual)
1478
+ def _shift_chunk(chunk: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
1479
+ out_i: list[int] = []
1480
+ out_dx: list[float] = []
1481
+ out_dy: list[float] = []
1482
+ out_cf: list[float] = []
1483
+
1484
+ src, owns = _ensure_source(source_obj, cache_items=0)
1485
+ try:
1486
+ for i in chunk.tolist():
1487
+ img = _get_frame(
1488
+ src, int(i),
1489
+ roi=roi_used,
1490
+ debayer=debayer,
1491
+ to_float01=True,
1492
+ force_rgb=bool(to_rgb),
1493
+ bayer_pattern=bpat,
1494
+ )
1495
+ cur_m = _to_mono01(img).astype(np.float32, copy=False)
1496
+
1497
+ coarse_dx = float(dx[int(i)])
1498
+ coarse_dy = float(dy[int(i)])
1499
+
1500
+ # Apply coarse shift FIRST (so APs line up without any searching)
1501
+ cur_m_g = _shift_image(cur_m, coarse_dx, coarse_dy)
1502
+
1503
+ if use_multiscale:
1504
+ s2, s1, s05 = _scaled_ap_sizes(ap_size)
1505
+
1506
+ def _one_scale(s_ap: int):
1507
+ rdx, rdy, resp = _ap_phase_shifts_per_ap(
1508
+ ref_m_full, cur_m_g,
1509
+ ap_centers=ap_centers,
1510
+ ap_size=s_ap,
1511
+ max_dim=max_dim,
1512
+ )
1513
+ cf = np.clip(resp.astype(np.float32, copy=False), 0.0, 1.0)
1514
+ keep = _reject_ap_outliers(rdx, rdy, cf, z=3.5)
1515
+ if not np.any(keep):
1516
+ return 0.0, 0.0, 0.25
1517
+ dx_r = float(np.median(rdx[keep]))
1518
+ dy_r = float(np.median(rdy[keep]))
1519
+ cf_r = float(np.median(cf[keep]))
1520
+ return dx_r, dy_r, cf_r
1521
+
1522
+ dx2, dy2, cf2 = _one_scale(s2)
1523
+ dx1, dy1, cf1 = _one_scale(s1)
1524
+ dx0, dy0, cf0 = _one_scale(s05)
1525
+
1526
+ w2 = max(1e-3, float(cf2)) * 1.25
1527
+ w1 = max(1e-3, float(cf1)) * 1.00
1528
+ w0 = max(1e-3, float(cf0)) * 0.85
1529
+ wsum = (w2 + w1 + w0)
1530
+
1531
+ dx_res = (w2 * dx2 + w1 * dx1 + w0 * dx0) / wsum
1532
+ dy_res = (w2 * dy2 + w1 * dy1 + w0 * dy0) / wsum
1533
+ cf_ap = float(np.clip((w2 * cf2 + w1 * cf1 + w0 * cf0) / wsum, 0.0, 1.0))
1534
+ else:
1535
+ rdx, rdy, resp = _ap_phase_shifts_per_ap(
1536
+ ref_m_full, cur_m_g,
1537
+ ap_centers=ap_centers,
1538
+ ap_size=ap_size,
1539
+ max_dim=max_dim,
1540
+ )
1541
+ cf = np.clip(resp.astype(np.float32, copy=False), 0.0, 1.0)
1542
+ keep = _reject_ap_outliers(rdx, rdy, cf, z=3.5)
1543
+ if np.any(keep):
1544
+ dx_res = float(np.median(rdx[keep]))
1545
+ dy_res = float(np.median(rdy[keep]))
1546
+ cf_ap = float(np.median(cf[keep]))
1547
+ else:
1548
+ dx_res, dy_res, cf_ap = 0.0, 0.0, 0.25
1549
+
1550
+ # Final = coarse + residual (residual is relative to coarse-shifted frame)
1551
+ # Final = coarse + residual (residual is relative to coarse-shifted frame)
1552
+ dx_i = float(coarse_dx + dx_res)
1553
+ dy_i = float(coarse_dy + dy_res)
1554
+
1555
+ # Final lock-in refinement: minimize (ref-cur)^2 on gradients in a tiny window
1556
+ # NOTE: pass *unshifted* cur_m with the current dx_i/dy_i estimate
1557
+ dxr, dyr, c_ssd = _refine_shift_ssd(
1558
+ ref_m_full, cur_m, dx_i, dy_i,
1559
+ radius=5, crop=0.80,
1560
+ bruteforce=bool(getattr(cfg, "ssd_refine_bruteforce", False)),
1561
+ )
1562
+ dx_i += float(dxr)
1563
+ dy_i += float(dyr)
1564
+
1565
+ # Confidence: combine coarse + AP, then optionally nudge with SSD
1566
+ cc = float(coarse_conf[int(i)]) if coarse_conf is not None else 0.5
1567
+ cf_i = float(np.clip(0.60 * cc + 0.40 * float(cf_ap), 0.0, 1.0))
1568
+ cf_i = float(np.clip(0.85 * cf_i + 0.15 * float(c_ssd), 0.05, 1.0))
1569
+
1570
+ out_i.append(int(i))
1571
+ out_dx.append(dx_i)
1572
+ out_dy.append(dy_i)
1573
+ out_cf.append(cf_i)
1574
+ finally:
1575
+ if owns:
1576
+ try:
1577
+ src.close()
1578
+ except Exception:
1579
+ pass
1580
+
1581
+ return (
1582
+ np.asarray(out_i, np.int32),
1583
+ np.asarray(out_dx, np.float32),
1584
+ np.asarray(out_dy, np.float32),
1585
+ np.asarray(out_cf, np.float32),
1586
+ )
1587
+
1588
+ else:
1589
+ # planetary: centroid tracking (same as viewer) for GLOBAL dx/dy/conf
1590
+ # APs are still computed and used later by stack_ser for local_warp residuals.
1591
+ tracker = PlanetaryTracker(
1592
+ smooth_sigma=float(getattr(cfg, "planet_smooth_sigma", smooth_sigma)),
1593
+ thresh_pct=float(getattr(cfg, "planet_thresh_pct", thresh_pct)),
1594
+ )
1595
+
1596
+ # IMPORTANT: reference center is computed from the SAME reference image that Analyze chose
1597
+ ref_cx, ref_cy, ref_cc = tracker.compute_center(ref_img)
1598
+ if ref_cc <= 0.0:
1599
+ # fallback: center of ROI
1600
+ mref = _to_mono01(ref_img)
1601
+ ref_cx = float(mref.shape[1] * 0.5)
1602
+ ref_cy = float(mref.shape[0] * 0.5)
1603
+
1604
+ ref_center = (float(ref_cx), float(ref_cy))
1605
+ ref_m_full = _to_mono01(ref_img).astype(np.float32, copy=False)
1606
+
1607
+ def _shift_chunk(chunk: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
1608
+ out_i: list[int] = []
1609
+ out_dx: list[float] = []
1610
+ out_dy: list[float] = []
1611
+ out_cf: list[float] = []
1612
+
1613
+ src, owns = _ensure_source(source_obj, cache_items=0)
1614
+ try:
1615
+ for i in chunk.tolist():
1616
+ img = _get_frame(
1617
+ src, int(i),
1618
+ roi=roi_used,
1619
+ debayer=debayer,
1620
+ to_float01=True,
1621
+ force_rgb=bool(to_rgb),
1622
+ bayer_pattern=bpat,
1623
+ )
1624
+
1625
+ dx_i, dy_i, cf_i = tracker.shift_to_ref(img, ref_center)
1626
+
1627
+ if float(cf_i) >= 0.25:
1628
+ cur_m = _to_mono01(img).astype(np.float32, copy=False)
1629
+ dxr, dyr, c_ssd = _refine_shift_ssd(
1630
+ ref_m_full, cur_m, dx_i, dy_i,
1631
+ radius=5, crop=0.80,
1632
+ bruteforce=bool(getattr(cfg, "ssd_refine_bruteforce", False)),
1633
+ )
1634
+
1635
+ dx_i = float(dx_i) + dxr
1636
+ dy_i = float(dy_i) + dyr
1637
+ cf_i = float(np.clip(0.85 * float(cf_i) + 0.15 * c_ssd, 0.05, 1.0))
1638
+ out_i.append(int(i))
1639
+ out_dx.append(float(dx_i))
1640
+ out_dy.append(float(dy_i))
1641
+ out_cf.append(float(cf_i))
1642
+ finally:
1643
+ if owns:
1644
+ try:
1645
+ src.close()
1646
+ except Exception:
1647
+ pass
1648
+
1649
+ return (
1650
+ np.asarray(out_i, np.int32),
1651
+ np.asarray(out_dx, np.float32),
1652
+ np.asarray(out_dy, np.float32),
1653
+ np.asarray(out_cf, np.float32),
1654
+ )
1655
+
1656
+
1657
+ done_ct = 0
1658
+ with ThreadPoolExecutor(max_workers=workers) as ex:
1659
+ futs = [ex.submit(_shift_chunk, c) for c in chunks2 if c.size > 0]
1660
+ for fut in as_completed(futs):
1661
+ ii, ddx, ddy, ccf = fut.result()
1662
+ dx[ii] = ddx
1663
+ dy[ii] = ddy
1664
+ conf[ii] = np.clip(ccf, 0.05, 1.0).astype(np.float32, copy=False)
1665
+
1666
+ done_ct += int(ii.size)
1667
+ if progress_cb:
1668
+ progress_cb(done_ct, n, "SSD Refine")
1669
+
1670
+ if cfg.track_mode == "surface":
1671
+ _print_surface_debug(dx=dx, dy=dy, conf=conf, coarse_conf=coarse_conf, floor=0.05, prefix="[SER][Surface]")
1672
+
1673
+ return AnalyzeResult(
1674
+ frames_total=n,
1675
+ roi_used=roi_used,
1676
+ track_mode=cfg.track_mode,
1677
+ quality=quality,
1678
+ dx=dx,
1679
+ dy=dy,
1680
+ conf=conf,
1681
+ order=order,
1682
+ ref_mode=ref_mode,
1683
+ ref_count=ref_count,
1684
+ ref_image=ref_img,
1685
+ ap_centers=ap_centers,
1686
+ ap_size=ap_size,
1687
+ ap_multiscale=use_multiscale,
1688
+ coarse_conf=coarse_conf,
1689
+ )
1690
+
1691
+
1692
+ def realign_ser(
1693
+ cfg: SERStackConfig,
1694
+ analysis: AnalyzeResult,
1695
+ *,
1696
+ debayer: bool = True,
1697
+ to_rgb: bool = False,
1698
+ max_dim: int = 512,
1699
+ progress_cb=None,
1700
+ workers: Optional[int] = None,
1701
+ ) -> AnalyzeResult:
1702
+ """
1703
+ Recompute dx/dy/conf only using analysis.ref_image and analysis.ap_centers.
1704
+ Keeps quality/order/ref_image unchanged.
1705
+
1706
+ Surface mode:
1707
+ - recompute coarse drift (ref-locked) on roi_track
1708
+ - refine via AP search+refine FOLLOWING coarse + outlier rejection
1709
+ """
1710
+ bpat = bayer_pattern or _cfg_bayer_pattern(cfg)
1711
+
1712
+ if analysis is None:
1713
+ raise ValueError("analysis is None")
1714
+ if analysis.ref_image is None:
1715
+ raise ValueError("analysis.ref_image is missing")
1716
+
1717
+ source_obj = _cfg_get_source(cfg)
1718
+ if not source_obj:
1719
+ raise ValueError("SERStackConfig.source/ser_path is empty")
1720
+
1721
+ n = int(analysis.frames_total)
1722
+ roi_used = analysis.roi_used
1723
+ ref_img = analysis.ref_image
1724
+
1725
+ if cfg.track_mode == "off" or cv2 is None:
1726
+ analysis.dx = np.zeros((n,), dtype=np.float32)
1727
+ analysis.dy = np.zeros((n,), dtype=np.float32)
1728
+ analysis.conf = np.ones((n,), dtype=np.float32)
1729
+ if hasattr(analysis, "coarse_conf"):
1730
+ analysis.coarse_conf = None
1731
+ return analysis
1732
+
1733
+ # Ensure AP centers exist
1734
+ ap_centers = getattr(analysis, "ap_centers", None)
1735
+ if ap_centers is None or np.asarray(ap_centers).size == 0:
1736
+ ap_centers = _autoplace_aps(
1737
+ ref_img,
1738
+ ap_size=int(getattr(cfg, "ap_size", 64)),
1739
+ ap_spacing=int(getattr(cfg, "ap_spacing", 48)),
1740
+ ap_min_mean=float(getattr(cfg, "ap_min_mean", 0.03)),
1741
+ )
1742
+ analysis.ap_centers = ap_centers
1743
+
1744
+ if workers is None:
1745
+ cpu = os.cpu_count() or 4
1746
+ workers = max(1, min(cpu, 48))
1747
+
1748
+ if cv2 is not None:
1749
+ try:
1750
+ cv2.setNumThreads(1)
1751
+ except Exception:
1752
+ pass
1753
+
1754
+ # Need meta for ROI expansion (surface tracking)
1755
+ src0, owns0 = _ensure_source(source_obj, cache_items=2)
1756
+ try:
1757
+ meta = src0.meta
1758
+ src_w = int(meta.width)
1759
+ src_h = int(meta.height)
1760
+ finally:
1761
+ if owns0:
1762
+ try:
1763
+ src0.close()
1764
+ except Exception:
1765
+ pass
1766
+
1767
+ def _surface_tracking_roi() -> Optional[Tuple[int, int, int, int]]:
1768
+ if roi_used is None:
1769
+ return None
1770
+ margin = int(getattr(cfg, "surface_track_margin", 256))
1771
+ x, y, w, h = [int(v) for v in roi_used]
1772
+ x0 = max(0, x - margin)
1773
+ y0 = max(0, y - margin)
1774
+ x1 = min(src_w, x + w + margin)
1775
+ y1 = min(src_h, y + h + margin)
1776
+ return _clamp_roi_in_bounds((x0, y0, x1 - x0, y1 - y0), src_w, src_h)
1777
+
1778
+ roi_track = _surface_tracking_roi() if cfg.track_mode == "surface" else roi_used
1779
+
1780
+ # ---- chunked refine ----
1781
+ idxs2 = np.arange(n, dtype=np.int32)
1782
+
1783
+ # More/smaller chunks => progress updates sooner (futures complete more frequently)
1784
+ chunk_factor = int(getattr(cfg, "progress_chunk_factor", 5)) # optional knob
1785
+ min_chunks = 5
1786
+ n_chunks2 = max(min_chunks, int(workers) * chunk_factor)
1787
+ n_chunks2 = max(1, min(int(n), n_chunks2))
1788
+
1789
+ chunks2 = np.array_split(idxs2, n_chunks2)
1790
+
1791
+ dx = np.zeros((n,), dtype=np.float32)
1792
+ dy = np.zeros((n,), dtype=np.float32)
1793
+ conf = np.ones((n,), dtype=np.float32)
1794
+
1795
+ ref_m = _to_mono01(ref_img).astype(np.float32, copy=False)
1796
+
1797
+ ap_size = int(getattr(cfg, "ap_size", 64) or 64)
1798
+ use_multiscale = bool(getattr(cfg, "ap_multiscale", False))
1799
+
1800
+ coarse_conf: Optional[np.ndarray] = None
1801
+ if cfg.track_mode == "surface":
1802
+ coarse_conf = np.zeros((n,), dtype=np.float32)
1803
+ if progress_cb:
1804
+ progress_cb(0, n, "Surface: coarse drift (ref-locked NCC+subpix)…")
1805
+
1806
+ dx_chain, dy_chain, cc_chain = _coarse_surface_ref_locked(
1807
+ source_obj,
1808
+ n=n,
1809
+ roi=roi_track,
1810
+ roi_used=roi_used, # ✅ NEW
1811
+ debayer=debayer,
1812
+ to_rgb=to_rgb,
1813
+ bayer_pattern=bpat,
1814
+ progress_cb=progress_cb,
1815
+ progress_every=25,
1816
+ down=2,
1817
+ template_size=256,
1818
+ search_radius=96,
1819
+ bandpass=True,
1820
+ workers=min(workers, 8), # coarse doesn’t need 48; 4–8 is usually ideal
1821
+ stride=16, # 8–32 typical
1822
+ )
1823
+
1824
+ dx[:] = dx_chain
1825
+ dy[:] = dy_chain
1826
+ coarse_conf[:] = cc_chain
1827
+
1828
+ if progress_cb:
1829
+ progress_cb(0, n, "SSD Refine")
1830
+
1831
+ if cfg.track_mode == "surface":
1832
+ def _shift_chunk(chunk: np.ndarray):
1833
+ out_i: list[int] = []
1834
+ out_dx: list[float] = []
1835
+ out_dy: list[float] = []
1836
+ out_cf: list[float] = []
1837
+ out_cc: list[float] = []
1838
+
1839
+ src, owns = _ensure_source(source_obj, cache_items=0)
1840
+ try:
1841
+ for i in chunk.tolist():
1842
+ img = _get_frame(
1843
+ src, int(i),
1844
+ roi=roi_used,
1845
+ debayer=debayer,
1846
+ to_float01=True,
1847
+ force_rgb=bool(to_rgb),
1848
+ bayer_pattern=bpat,
1849
+ )
1850
+ cur_m = _to_mono01(img).astype(np.float32, copy=False)
1851
+
1852
+ coarse_dx = float(dx[int(i)])
1853
+ coarse_dy = float(dy[int(i)])
1854
+ cc = float(coarse_conf[int(i)]) if coarse_conf is not None else 0.5
1855
+
1856
+ # Apply coarse shift first
1857
+ cur_m_g = _shift_image(cur_m, coarse_dx, coarse_dy)
1858
+
1859
+ if use_multiscale:
1860
+ s2, s1, s05 = _scaled_ap_sizes(ap_size)
1861
+
1862
+ def _one_scale(s_ap: int):
1863
+ rdx, rdy, resp = _ap_phase_shifts_per_ap(
1864
+ ref_m, cur_m_g,
1865
+ ap_centers=ap_centers,
1866
+ ap_size=s_ap,
1867
+ max_dim=max_dim,
1868
+ )
1869
+ cf = np.clip(resp.astype(np.float32, copy=False), 0.0, 1.0)
1870
+ keep = _reject_ap_outliers(rdx, rdy, cf, z=3.5)
1871
+ if not np.any(keep):
1872
+ return 0.0, 0.0, 0.25
1873
+ return (
1874
+ float(np.median(rdx[keep])),
1875
+ float(np.median(rdy[keep])),
1876
+ float(np.median(cf[keep])),
1877
+ )
1878
+
1879
+ dx2, dy2, cf2 = _one_scale(s2)
1880
+ dx1, dy1, cf1 = _one_scale(s1)
1881
+ dx0, dy0, cf0 = _one_scale(s05)
1882
+
1883
+ w2 = max(1e-3, float(cf2)) * 1.25
1884
+ w1 = max(1e-3, float(cf1)) * 1.00
1885
+ w0 = max(1e-3, float(cf0)) * 0.85
1886
+ wsum = (w2 + w1 + w0)
1887
+
1888
+ dx_res = (w2 * dx2 + w1 * dx1 + w0 * dx0) / wsum
1889
+ dy_res = (w2 * dy2 + w1 * dy1 + w0 * dy0) / wsum
1890
+ cf_ap = float(np.clip((w2 * cf2 + w1 * cf1 + w0 * cf0) / wsum, 0.0, 1.0))
1891
+ else:
1892
+ rdx, rdy, resp = _ap_phase_shifts_per_ap(
1893
+ ref_m, cur_m_g,
1894
+ ap_centers=ap_centers,
1895
+ ap_size=ap_size,
1896
+ max_dim=max_dim,
1897
+ )
1898
+ cf = np.clip(resp.astype(np.float32, copy=False), 0.0, 1.0)
1899
+ keep = _reject_ap_outliers(rdx, rdy, cf, z=3.5)
1900
+ if np.any(keep):
1901
+ dx_res = float(np.median(rdx[keep]))
1902
+ dy_res = float(np.median(rdy[keep]))
1903
+ cf_ap = float(np.median(cf[keep]))
1904
+ else:
1905
+ dx_res, dy_res, cf_ap = 0.0, 0.0, 0.25
1906
+
1907
+ # Final = coarse + residual (residual is relative to coarse-shifted frame)
1908
+ dx_i = float(coarse_dx + dx_res)
1909
+ dy_i = float(coarse_dy + dy_res)
1910
+
1911
+ # Final lock-in refinement: minimize (ref-cur)^2 on gradients in a tiny window
1912
+ # NOTE: pass *unshifted* cur_m with the current dx_i/dy_i estimate
1913
+ dxr, dyr, c_ssd = _refine_shift_ssd(ref_m, cur_m, dx_i, dy_i, radius=5, crop=0.80, bruteforce=bool(getattr(cfg, "ssd_refine_bruteforce", False)))
1914
+ dx_i += float(dxr)
1915
+ dy_i += float(dyr)
1916
+
1917
+ # Confidence: combine coarse + AP, then optionally nudge with SSD
1918
+ cc = float(coarse_conf[int(i)]) if coarse_conf is not None else 0.5
1919
+ cf_i = float(np.clip(0.60 * cc + 0.40 * float(cf_ap), 0.0, 1.0))
1920
+ cf_i = float(np.clip(0.85 * cf_i + 0.15 * float(c_ssd), 0.05, 1.0))
1921
+
1922
+
1923
+ out_i.append(int(i))
1924
+ out_dx.append(dx_i)
1925
+ out_dy.append(dy_i)
1926
+ out_cf.append(cf_i)
1927
+ out_cc.append(float(cc))
1928
+ finally:
1929
+ if owns:
1930
+ try:
1931
+ src.close()
1932
+ except Exception:
1933
+ pass
1934
+
1935
+ return (
1936
+ np.asarray(out_i, np.int32),
1937
+ np.asarray(out_dx, np.float32),
1938
+ np.asarray(out_dy, np.float32),
1939
+ np.asarray(out_cf, np.float32),
1940
+ np.asarray(out_cc, np.float32),
1941
+ )
1942
+
1943
+ else:
1944
+ # planetary: centroid tracking (same as viewer)
1945
+ tracker = PlanetaryTracker(
1946
+ smooth_sigma=float(getattr(cfg, "planet_smooth_sigma", 1.5)),
1947
+ thresh_pct=float(getattr(cfg, "planet_thresh_pct", 92.0)),
1948
+ )
1949
+
1950
+ # Reference center comes from analysis.ref_image (same anchor as analyze_ser)
1951
+ ref_cx, ref_cy, ref_cc = tracker.compute_center(ref_img)
1952
+ if ref_cc <= 0.0:
1953
+ mref = _to_mono01(ref_img)
1954
+ ref_cx = float(mref.shape[1] * 0.5)
1955
+ ref_cy = float(mref.shape[0] * 0.5)
1956
+
1957
+ ref_center = (float(ref_cx), float(ref_cy))
1958
+ ref_m_full = _to_mono01(ref_img).astype(np.float32, copy=False)
1959
+
1960
+ def _shift_chunk(chunk: np.ndarray):
1961
+ out_i: list[int] = []
1962
+ out_dx: list[float] = []
1963
+ out_dy: list[float] = []
1964
+ out_cf: list[float] = []
1965
+
1966
+ src, owns = _ensure_source(source_obj, cache_items=0)
1967
+ try:
1968
+ for i in chunk.tolist():
1969
+ img = _get_frame(
1970
+ src, int(i),
1971
+ roi=roi_used,
1972
+ debayer=debayer,
1973
+ to_float01=True,
1974
+ force_rgb=bool(to_rgb),
1975
+ bayer_pattern=bpat,
1976
+ )
1977
+
1978
+ dx_i, dy_i, cf_i = tracker.shift_to_ref(img, ref_center)
1979
+
1980
+ if float(cf_i) >= 0.25:
1981
+ cur_m = _to_mono01(img).astype(np.float32, copy=False)
1982
+ dxr, dyr, c_ssd = _refine_shift_ssd(ref_m_full, cur_m, float(dx_i), float(dy_i), radius=2, crop=0.80, bruteforce=bool(getattr(cfg, "ssd_refine_bruteforce", False)))
1983
+ dx_i = float(dx_i) + dxr
1984
+ dy_i = float(dy_i) + dyr
1985
+ cf_i = float(np.clip(0.85 * float(cf_i) + 0.15 * c_ssd, 0.05, 1.0))
1986
+ out_i.append(int(i))
1987
+ out_dx.append(float(dx_i))
1988
+ out_dy.append(float(dy_i))
1989
+ out_cf.append(float(cf_i))
1990
+
1991
+ finally:
1992
+ if owns:
1993
+ try:
1994
+ src.close()
1995
+ except Exception:
1996
+ pass
1997
+
1998
+ return (
1999
+ np.asarray(out_i, np.int32),
2000
+ np.asarray(out_dx, np.float32),
2001
+ np.asarray(out_dy, np.float32),
2002
+ np.asarray(out_cf, np.float32),
2003
+ )
2004
+
2005
+
2006
+ done_ct = 0
2007
+ with ThreadPoolExecutor(max_workers=workers) as ex:
2008
+ futs = [ex.submit(_shift_chunk, c) for c in chunks2 if c.size > 0]
2009
+ for fut in as_completed(futs):
2010
+ if cfg.track_mode == "surface":
2011
+ ii, ddx, ddy, ccf, ccc = fut.result()
2012
+ if coarse_conf is not None:
2013
+ coarse_conf[ii] = ccc
2014
+ else:
2015
+ ii, ddx, ddy, ccf = fut.result()
2016
+
2017
+ dx[ii] = ddx
2018
+ dy[ii] = ddy
2019
+ conf[ii] = np.clip(ccf, 0.05, 1.0).astype(np.float32, copy=False)
2020
+
2021
+ done_ct += int(ii.size)
2022
+ if progress_cb:
2023
+ progress_cb(done_ct, n, "SSD Refine")
2024
+
2025
+ analysis.dx = dx
2026
+ analysis.dy = dy
2027
+ analysis.conf = conf
2028
+ if hasattr(analysis, "coarse_conf"):
2029
+ analysis.coarse_conf = coarse_conf
2030
+
2031
+ if cfg.track_mode == "surface":
2032
+ _print_surface_debug(dx=dx, dy=dy, conf=conf, coarse_conf=coarse_conf, floor=0.05, prefix="[SER][Surface][realign]")
2033
+
2034
+ return analysis
2035
+
2036
+ def _autoplace_aps(ref_img01: np.ndarray, ap_size: int, ap_spacing: int, ap_min_mean: float) -> np.ndarray:
2037
+ """
2038
+ Return AP centers as int32 array of shape (M,2) with columns (cx, cy) in ROI coords.
2039
+ We grid-scan by spacing and keep patches whose mean brightness exceeds ap_min_mean.
2040
+ """
2041
+ m = _to_mono01(ref_img01).astype(np.float32, copy=False)
2042
+ H, W = m.shape[:2]
2043
+ s = int(max(16, ap_size))
2044
+ step = int(max(4, ap_spacing))
2045
+
2046
+ half = s // 2
2047
+ xs = list(range(half, max(half + 1, W - half), step))
2048
+ ys = list(range(half, max(half + 1, H - half), step))
2049
+
2050
+ pts = []
2051
+ for cy in ys:
2052
+ y0 = cy - half
2053
+ y1 = y0 + s
2054
+ if y0 < 0 or y1 > H:
2055
+ continue
2056
+ for cx in xs:
2057
+ x0 = cx - half
2058
+ x1 = x0 + s
2059
+ if x0 < 0 or x1 > W:
2060
+ continue
2061
+ patch = m[y0:y1, x0:x1]
2062
+ if float(patch.mean()) >= float(ap_min_mean):
2063
+ pts.append((cx, cy))
2064
+
2065
+ if not pts:
2066
+ # absolute fallback: a single center point (behaves like single-point)
2067
+ pts = [(W // 2, H // 2)]
2068
+
2069
+ return np.asarray(pts, dtype=np.int32)
2070
+
2071
+ def _scaled_ap_sizes(base: int) -> tuple[int, int, int]:
2072
+ b = int(base)
2073
+ s2 = int(round(b * 2.0))
2074
+ s1 = int(round(b * 1.0))
2075
+ s05 = int(round(b * 0.5))
2076
+ # clamp to sane limits
2077
+ s2 = max(16, min(256, s2))
2078
+ s1 = max(16, min(256, s1))
2079
+ s05 = max(16, min(256, s05))
2080
+ return s2, s1, s05
2081
+
2082
+ def _dense_field_from_ap_shifts(
2083
+ H: int, W: int,
2084
+ ap_centers: np.ndarray, # (M,2)
2085
+ ap_dx: np.ndarray, # (M,)
2086
+ ap_dy: np.ndarray, # (M,)
2087
+ ap_cf: np.ndarray, # (M,)
2088
+ *,
2089
+ grid: int = 32, # coarse grid resolution (32 or 48 are good)
2090
+ power: float = 2.0,
2091
+ conf_floor: float = 0.15,
2092
+ radius: float | None = None, # optional clamp in pixels (ROI coords)
2093
+ ) -> tuple[np.ndarray, np.ndarray]:
2094
+ """
2095
+ Returns dense (dx_field, dy_field) as float32 arrays (H,W) in ROI pixels.
2096
+ Computed on coarse grid then upsampled.
2097
+ """
2098
+ # coarse grid points
2099
+ gh = max(4, int(grid))
2100
+ gw = max(4, int(round(grid * (W / max(1, H)))))
2101
+
2102
+ ys = np.linspace(0, H - 1, gh, dtype=np.float32)
2103
+ xs = np.linspace(0, W - 1, gw, dtype=np.float32)
2104
+ gx, gy = np.meshgrid(xs, ys) # (gh,gw)
2105
+
2106
+ pts = ap_centers.astype(np.float32)
2107
+ px = pts[:, 0].reshape(-1, 1, 1) # (M,1,1)
2108
+ py = pts[:, 1].reshape(-1, 1, 1) # (M,1,1)
2109
+
2110
+ cf = np.maximum(ap_cf.astype(np.float32), 0.0)
2111
+ good = cf >= float(conf_floor)
2112
+
2113
+ if not np.any(good):
2114
+ dxg = np.zeros((gh, gw), np.float32)
2115
+ dyg = np.zeros((gh, gw), np.float32)
2116
+ else:
2117
+ px = px[good]
2118
+ py = py[good]
2119
+ dx = ap_dx[good].astype(np.float32).reshape(-1, 1, 1)
2120
+ dy = ap_dy[good].astype(np.float32).reshape(-1, 1, 1)
2121
+ cw = cf[good].astype(np.float32).reshape(-1, 1, 1)
2122
+
2123
+ dxp = px - gx[None, :, :] # (M,gh,gw)
2124
+ dyp = py - gy[None, :, :] # (M,gh,gw)
2125
+ d2 = dxp * dxp + dyp * dyp # (M,gh,gw)
2126
+
2127
+ if radius is not None:
2128
+ r2 = float(radius) * float(radius)
2129
+ far = d2 > r2
2130
+ else:
2131
+ far = None
2132
+
2133
+ w = 1.0 / np.maximum(d2, 1.0) ** (power * 0.5)
2134
+ w *= cw
2135
+
2136
+ if far is not None:
2137
+ w = np.where(far, 0.0, w)
2138
+
2139
+ wsum = np.sum(w, axis=0) # (gh,gw)
2140
+
2141
+ dxg = np.sum(w * dx, axis=0) / np.maximum(wsum, 1e-6)
2142
+ dyg = np.sum(w * dy, axis=0) / np.maximum(wsum, 1e-6)
2143
+
2144
+
2145
+ # upsample to full res
2146
+ dx_field = cv2.resize(dxg, (W, H), interpolation=cv2.INTER_CUBIC).astype(np.float32, copy=False)
2147
+ dy_field = cv2.resize(dyg, (W, H), interpolation=cv2.INTER_CUBIC).astype(np.float32, copy=False)
2148
+ return dx_field, dy_field
2149
+
2150
+ def _warp_by_dense_field(img01: np.ndarray, dx_field: np.ndarray, dy_field: np.ndarray) -> np.ndarray:
2151
+ """
2152
+ img01 (H,W) or (H,W,3)
2153
+ dx_field/dy_field are (H,W) in pixels: shifting cur by (dx,dy) aligns to ref.
2154
+ """
2155
+ H, W = dx_field.shape
2156
+ # remap wants map_x/map_y = source sampling coordinates
2157
+ # If we want output aligned-to-ref, we sample from cur at (x - dx, y - dy)
2158
+ xs, ys = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32))
2159
+ map_x = xs - dx_field
2160
+ map_y = ys - dy_field
2161
+
2162
+ if img01.ndim == 2:
2163
+ return cv2.remap(img01, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
2164
+ else:
2165
+ return cv2.remap(img01, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
2166
+
2167
+ def _ap_phase_shift(
2168
+ ref_m: np.ndarray,
2169
+ cur_m: np.ndarray,
2170
+ ap_centers: np.ndarray,
2171
+ ap_size: int,
2172
+ max_dim: int,
2173
+ ) -> tuple[float, float, float]:
2174
+ """
2175
+ Compute a robust global shift from multiple local AP shifts.
2176
+ Returns (dx, dy, conf) in ROI pixel units.
2177
+ conf is median of per-AP phase correlation responses.
2178
+ """
2179
+ s = int(max(16, ap_size))
2180
+ half = s // 2
2181
+
2182
+ H, W = ref_m.shape[:2]
2183
+ dxs = []
2184
+ dys = []
2185
+ resps = []
2186
+
2187
+ # downsample reference patches once per AP? (fast enough as-is; M is usually modest)
2188
+ for (cx, cy) in ap_centers.tolist():
2189
+ x0 = cx - half
2190
+ y0 = cy - half
2191
+ x1 = x0 + s
2192
+ y1 = y0 + s
2193
+ if x0 < 0 or y0 < 0 or x1 > W or y1 > H:
2194
+ continue
2195
+
2196
+ ref_patch = ref_m[y0:y1, x0:x1]
2197
+ cur_patch = cur_m[y0:y1, x0:x1]
2198
+
2199
+ rp = _downsample_mono01(ref_patch, max_dim=max_dim)
2200
+ cp = _downsample_mono01(cur_patch, max_dim=max_dim)
2201
+
2202
+ if rp.shape != cp.shape:
2203
+ cp = cv2.resize(cp, (rp.shape[1], rp.shape[0]), interpolation=cv2.INTER_AREA)
2204
+
2205
+ sdx, sdy, resp = _phase_corr_shift(rp, cp)
2206
+
2207
+ # scale back to ROI pixels (patch pixels -> ROI pixels)
2208
+ sx = float(s) / float(rp.shape[1])
2209
+ sy = float(s) / float(rp.shape[0])
2210
+
2211
+ dxs.append(float(sdx * sx))
2212
+ dys.append(float(sdy * sy))
2213
+ resps.append(float(resp))
2214
+
2215
+ if not dxs:
2216
+ return 0.0, 0.0, 0.5
2217
+
2218
+ dx_med = float(np.median(np.asarray(dxs, np.float32)))
2219
+ dy_med = float(np.median(np.asarray(dys, np.float32)))
2220
+ conf = float(np.median(np.asarray(resps, np.float32)))
2221
+
2222
+ return dx_med, dy_med, conf
2223
+
2224
+ def _ap_phase_shifts_per_ap(
2225
+ ref_m: np.ndarray,
2226
+ cur_m: np.ndarray,
2227
+ ap_centers: np.ndarray,
2228
+ ap_size: int,
2229
+ max_dim: int,
2230
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
2231
+ """
2232
+ Per-AP phase correlation shifts (NO SEARCH).
2233
+ Returns arrays (ap_dx, ap_dy, ap_resp) in ROI pixels, where shifting cur by (dx,dy)
2234
+ aligns it to ref for each AP.
2235
+ """
2236
+ s = int(max(16, ap_size))
2237
+ half = s // 2
2238
+
2239
+ H, W = ref_m.shape[:2]
2240
+ M = int(ap_centers.shape[0])
2241
+
2242
+ ap_dx = np.zeros((M,), np.float32)
2243
+ ap_dy = np.zeros((M,), np.float32)
2244
+ ap_resp = np.zeros((M,), np.float32)
2245
+
2246
+ if cv2 is None or M == 0:
2247
+ ap_resp[:] = 0.5
2248
+ return ap_dx, ap_dy, ap_resp
2249
+
2250
+ for j, (cx, cy) in enumerate(ap_centers.tolist()):
2251
+ x0 = int(cx - half)
2252
+ y0 = int(cy - half)
2253
+ x1 = x0 + s
2254
+ y1 = y0 + s
2255
+ if x0 < 0 or y0 < 0 or x1 > W or y1 > H:
2256
+ ap_resp[j] = 0.0
2257
+ continue
2258
+
2259
+ ref_patch = ref_m[y0:y1, x0:x1]
2260
+ cur_patch = cur_m[y0:y1, x0:x1]
2261
+
2262
+ rp = _downsample_mono01(ref_patch, max_dim=max_dim)
2263
+ cp = _downsample_mono01(cur_patch, max_dim=max_dim)
2264
+
2265
+ if rp.shape != cp.shape and cv2 is not None:
2266
+ cp = cv2.resize(cp, (rp.shape[1], rp.shape[0]), interpolation=cv2.INTER_AREA)
2267
+
2268
+ sdx, sdy, resp = _phase_corr_shift(rp, cp)
2269
+
2270
+ # scale to ROI pixels
2271
+ sx = float(s) / float(rp.shape[1])
2272
+ sy = float(s) / float(rp.shape[0])
2273
+
2274
+ ap_dx[j] = float(sdx * sx)
2275
+ ap_dy[j] = float(sdy * sy)
2276
+ ap_resp[j] = float(resp)
2277
+
2278
+ return ap_dx, ap_dy, ap_resp
2279
+
2280
+
2281
+ def _ap_phase_shift_multiscale(
2282
+ ref_m: np.ndarray,
2283
+ cur_m: np.ndarray,
2284
+ ap_centers: np.ndarray,
2285
+ base_ap_size: int,
2286
+ max_dim: int,
2287
+ ) -> tuple[float, float, float]:
2288
+ """
2289
+ Multi-scale AP shift:
2290
+ - compute shifts at 2×, 1×, ½× AP sizes using same centers
2291
+ - combine using confidence weights (favoring coarser slightly)
2292
+ Returns (dx, dy, conf) in ROI pixels.
2293
+ """
2294
+ s2, s1, s05 = _scaled_ap_sizes(base_ap_size)
2295
+
2296
+ dx2, dy2, cf2 = _ap_phase_shift(ref_m, cur_m, ap_centers, s2, max_dim)
2297
+ dx1, dy1, cf1 = _ap_phase_shift(ref_m, cur_m, ap_centers, s1, max_dim)
2298
+ dx0, dy0, cf0 = _ap_phase_shift(ref_m, cur_m, ap_centers, s05, max_dim)
2299
+
2300
+ # weights: confidence * slight preference for larger scale (stability)
2301
+ w2 = max(1e-3, float(cf2)) * 1.25
2302
+ w1 = max(1e-3, float(cf1)) * 1.00
2303
+ w0 = max(1e-3, float(cf0)) * 0.85
2304
+
2305
+ wsum = (w2 + w1 + w0)
2306
+ dx = (w2 * dx2 + w1 * dx1 + w0 * dx0) / wsum
2307
+ dy = (w2 * dy2 + w1 * dy1 + w0 * dy0) / wsum
2308
+ conf = float(np.clip((w2 * cf2 + w1 * cf1 + w0 * cf0) / wsum, 0.0, 1.0))
2309
+
2310
+ return float(dx), float(dy), float(conf)