setiastrosuitepro 1.6.7__py3-none-any.whl → 1.7.0__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.

Potentially problematic release.


This version of setiastrosuitepro might be problematic. Click here for more details.

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