setiastrosuitepro 1.6.2.post1__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 (367) hide show
  1. setiastro/__init__.py +2 -0
  2. setiastro/data/SASP_data.fits +0 -0
  3. setiastro/data/catalogs/List_of_Galaxies_with_Distances_Gly.csv +488 -0
  4. setiastro/data/catalogs/astrobin_filters.csv +890 -0
  5. setiastro/data/catalogs/astrobin_filters_page1_local.csv +51 -0
  6. setiastro/data/catalogs/cali2.csv +63 -0
  7. setiastro/data/catalogs/cali2color.csv +65 -0
  8. setiastro/data/catalogs/celestial_catalog - original.csv +16471 -0
  9. setiastro/data/catalogs/celestial_catalog.csv +24031 -0
  10. setiastro/data/catalogs/detected_stars.csv +24784 -0
  11. setiastro/data/catalogs/fits_header_data.csv +46 -0
  12. setiastro/data/catalogs/test.csv +8 -0
  13. setiastro/data/catalogs/updated_celestial_catalog.csv +16471 -0
  14. setiastro/images/Astro_Spikes.png +0 -0
  15. setiastro/images/Background_startup.jpg +0 -0
  16. setiastro/images/HRDiagram.png +0 -0
  17. setiastro/images/LExtract.png +0 -0
  18. setiastro/images/LInsert.png +0 -0
  19. setiastro/images/Oxygenation-atm-2.svg.png +0 -0
  20. setiastro/images/RGB080604.png +0 -0
  21. setiastro/images/abeicon.png +0 -0
  22. setiastro/images/aberration.png +0 -0
  23. setiastro/images/andromedatry.png +0 -0
  24. setiastro/images/andromedatry_satellited.png +0 -0
  25. setiastro/images/annotated.png +0 -0
  26. setiastro/images/aperture.png +0 -0
  27. setiastro/images/astrosuite.ico +0 -0
  28. setiastro/images/astrosuite.png +0 -0
  29. setiastro/images/astrosuitepro.icns +0 -0
  30. setiastro/images/astrosuitepro.ico +0 -0
  31. setiastro/images/astrosuitepro.png +0 -0
  32. setiastro/images/background.png +0 -0
  33. setiastro/images/background2.png +0 -0
  34. setiastro/images/benchmark.png +0 -0
  35. setiastro/images/big_moon_stabilizer_timeline.png +0 -0
  36. setiastro/images/big_moon_stabilizer_timeline_clean.png +0 -0
  37. setiastro/images/blaster.png +0 -0
  38. setiastro/images/blink.png +0 -0
  39. setiastro/images/clahe.png +0 -0
  40. setiastro/images/collage.png +0 -0
  41. setiastro/images/colorwheel.png +0 -0
  42. setiastro/images/contsub.png +0 -0
  43. setiastro/images/convo.png +0 -0
  44. setiastro/images/copyslot.png +0 -0
  45. setiastro/images/cosmic.png +0 -0
  46. setiastro/images/cosmicsat.png +0 -0
  47. setiastro/images/crop1.png +0 -0
  48. setiastro/images/cropicon.png +0 -0
  49. setiastro/images/curves.png +0 -0
  50. setiastro/images/cvs.png +0 -0
  51. setiastro/images/debayer.png +0 -0
  52. setiastro/images/denoise_cnn_custom.png +0 -0
  53. setiastro/images/denoise_cnn_graph.png +0 -0
  54. setiastro/images/disk.png +0 -0
  55. setiastro/images/dse.png +0 -0
  56. setiastro/images/exoicon.png +0 -0
  57. setiastro/images/eye.png +0 -0
  58. setiastro/images/fliphorizontal.png +0 -0
  59. setiastro/images/flipvertical.png +0 -0
  60. setiastro/images/font.png +0 -0
  61. setiastro/images/freqsep.png +0 -0
  62. setiastro/images/functionbundle.png +0 -0
  63. setiastro/images/graxpert.png +0 -0
  64. setiastro/images/green.png +0 -0
  65. setiastro/images/gridicon.png +0 -0
  66. setiastro/images/halo.png +0 -0
  67. setiastro/images/hdr.png +0 -0
  68. setiastro/images/histogram.png +0 -0
  69. setiastro/images/hubble.png +0 -0
  70. setiastro/images/imagecombine.png +0 -0
  71. setiastro/images/invert.png +0 -0
  72. setiastro/images/isophote.png +0 -0
  73. setiastro/images/isophote_demo_figure.png +0 -0
  74. setiastro/images/isophote_demo_image.png +0 -0
  75. setiastro/images/isophote_demo_model.png +0 -0
  76. setiastro/images/isophote_demo_residual.png +0 -0
  77. setiastro/images/jwstpupil.png +0 -0
  78. setiastro/images/linearfit.png +0 -0
  79. setiastro/images/livestacking.png +0 -0
  80. setiastro/images/mask.png +0 -0
  81. setiastro/images/maskapply.png +0 -0
  82. setiastro/images/maskcreate.png +0 -0
  83. setiastro/images/maskremove.png +0 -0
  84. setiastro/images/morpho.png +0 -0
  85. setiastro/images/mosaic.png +0 -0
  86. setiastro/images/multiscale_decomp.png +0 -0
  87. setiastro/images/nbtorgb.png +0 -0
  88. setiastro/images/neutral.png +0 -0
  89. setiastro/images/nuke.png +0 -0
  90. setiastro/images/openfile.png +0 -0
  91. setiastro/images/pedestal.png +0 -0
  92. setiastro/images/pen.png +0 -0
  93. setiastro/images/pixelmath.png +0 -0
  94. setiastro/images/platesolve.png +0 -0
  95. setiastro/images/ppp.png +0 -0
  96. setiastro/images/pro.png +0 -0
  97. setiastro/images/project.png +0 -0
  98. setiastro/images/psf.png +0 -0
  99. setiastro/images/redo.png +0 -0
  100. setiastro/images/redoicon.png +0 -0
  101. setiastro/images/rescale.png +0 -0
  102. setiastro/images/rgbalign.png +0 -0
  103. setiastro/images/rgbcombo.png +0 -0
  104. setiastro/images/rgbextract.png +0 -0
  105. setiastro/images/rotate180.png +0 -0
  106. setiastro/images/rotateclockwise.png +0 -0
  107. setiastro/images/rotatecounterclockwise.png +0 -0
  108. setiastro/images/satellite.png +0 -0
  109. setiastro/images/script.png +0 -0
  110. setiastro/images/selectivecolor.png +0 -0
  111. setiastro/images/simbad.png +0 -0
  112. setiastro/images/slot0.png +0 -0
  113. setiastro/images/slot1.png +0 -0
  114. setiastro/images/slot2.png +0 -0
  115. setiastro/images/slot3.png +0 -0
  116. setiastro/images/slot4.png +0 -0
  117. setiastro/images/slot5.png +0 -0
  118. setiastro/images/slot6.png +0 -0
  119. setiastro/images/slot7.png +0 -0
  120. setiastro/images/slot8.png +0 -0
  121. setiastro/images/slot9.png +0 -0
  122. setiastro/images/spcc.png +0 -0
  123. setiastro/images/spin_precession_vs_lunar_distance.png +0 -0
  124. setiastro/images/spinner.gif +0 -0
  125. setiastro/images/stacking.png +0 -0
  126. setiastro/images/staradd.png +0 -0
  127. setiastro/images/staralign.png +0 -0
  128. setiastro/images/starnet.png +0 -0
  129. setiastro/images/starregistration.png +0 -0
  130. setiastro/images/starspike.png +0 -0
  131. setiastro/images/starstretch.png +0 -0
  132. setiastro/images/statstretch.png +0 -0
  133. setiastro/images/supernova.png +0 -0
  134. setiastro/images/uhs.png +0 -0
  135. setiastro/images/undoicon.png +0 -0
  136. setiastro/images/upscale.png +0 -0
  137. setiastro/images/viewbundle.png +0 -0
  138. setiastro/images/whitebalance.png +0 -0
  139. setiastro/images/wimi_icon_256x256.png +0 -0
  140. setiastro/images/wimilogo.png +0 -0
  141. setiastro/images/wims.png +0 -0
  142. setiastro/images/wrench_icon.png +0 -0
  143. setiastro/images/xisfliberator.png +0 -0
  144. setiastro/qml/ResourceMonitor.qml +126 -0
  145. setiastro/saspro/__init__.py +20 -0
  146. setiastro/saspro/__main__.py +945 -0
  147. setiastro/saspro/_generated/__init__.py +7 -0
  148. setiastro/saspro/_generated/build_info.py +3 -0
  149. setiastro/saspro/abe.py +1346 -0
  150. setiastro/saspro/abe_preset.py +196 -0
  151. setiastro/saspro/aberration_ai.py +694 -0
  152. setiastro/saspro/aberration_ai_preset.py +224 -0
  153. setiastro/saspro/accel_installer.py +218 -0
  154. setiastro/saspro/accel_workers.py +30 -0
  155. setiastro/saspro/add_stars.py +624 -0
  156. setiastro/saspro/astrobin_exporter.py +1010 -0
  157. setiastro/saspro/astrospike.py +153 -0
  158. setiastro/saspro/astrospike_python.py +1841 -0
  159. setiastro/saspro/autostretch.py +198 -0
  160. setiastro/saspro/backgroundneutral.py +602 -0
  161. setiastro/saspro/batch_convert.py +328 -0
  162. setiastro/saspro/batch_renamer.py +522 -0
  163. setiastro/saspro/blemish_blaster.py +491 -0
  164. setiastro/saspro/blink_comparator_pro.py +2926 -0
  165. setiastro/saspro/bundles.py +61 -0
  166. setiastro/saspro/bundles_dock.py +114 -0
  167. setiastro/saspro/cheat_sheet.py +213 -0
  168. setiastro/saspro/clahe.py +368 -0
  169. setiastro/saspro/comet_stacking.py +1442 -0
  170. setiastro/saspro/common_tr.py +107 -0
  171. setiastro/saspro/config.py +38 -0
  172. setiastro/saspro/config_bootstrap.py +40 -0
  173. setiastro/saspro/config_manager.py +316 -0
  174. setiastro/saspro/continuum_subtract.py +1617 -0
  175. setiastro/saspro/convo.py +1400 -0
  176. setiastro/saspro/convo_preset.py +414 -0
  177. setiastro/saspro/copyastro.py +190 -0
  178. setiastro/saspro/cosmicclarity.py +1589 -0
  179. setiastro/saspro/cosmicclarity_preset.py +407 -0
  180. setiastro/saspro/crop_dialog_pro.py +973 -0
  181. setiastro/saspro/crop_preset.py +189 -0
  182. setiastro/saspro/curve_editor_pro.py +2562 -0
  183. setiastro/saspro/curves_preset.py +375 -0
  184. setiastro/saspro/debayer.py +673 -0
  185. setiastro/saspro/debug_utils.py +29 -0
  186. setiastro/saspro/dnd_mime.py +35 -0
  187. setiastro/saspro/doc_manager.py +2664 -0
  188. setiastro/saspro/exoplanet_detector.py +2166 -0
  189. setiastro/saspro/file_utils.py +284 -0
  190. setiastro/saspro/fitsmodifier.py +748 -0
  191. setiastro/saspro/fix_bom.py +32 -0
  192. setiastro/saspro/free_torch_memory.py +48 -0
  193. setiastro/saspro/frequency_separation.py +1349 -0
  194. setiastro/saspro/function_bundle.py +1596 -0
  195. setiastro/saspro/generate_translations.py +3092 -0
  196. setiastro/saspro/ghs_dialog_pro.py +663 -0
  197. setiastro/saspro/ghs_preset.py +284 -0
  198. setiastro/saspro/graxpert.py +637 -0
  199. setiastro/saspro/graxpert_preset.py +287 -0
  200. setiastro/saspro/gui/__init__.py +0 -0
  201. setiastro/saspro/gui/main_window.py +8810 -0
  202. setiastro/saspro/gui/mixins/__init__.py +33 -0
  203. setiastro/saspro/gui/mixins/dock_mixin.py +362 -0
  204. setiastro/saspro/gui/mixins/file_mixin.py +450 -0
  205. setiastro/saspro/gui/mixins/geometry_mixin.py +403 -0
  206. setiastro/saspro/gui/mixins/header_mixin.py +441 -0
  207. setiastro/saspro/gui/mixins/mask_mixin.py +421 -0
  208. setiastro/saspro/gui/mixins/menu_mixin.py +389 -0
  209. setiastro/saspro/gui/mixins/theme_mixin.py +367 -0
  210. setiastro/saspro/gui/mixins/toolbar_mixin.py +1457 -0
  211. setiastro/saspro/gui/mixins/update_mixin.py +309 -0
  212. setiastro/saspro/gui/mixins/view_mixin.py +435 -0
  213. setiastro/saspro/gui/statistics_dialog.py +47 -0
  214. setiastro/saspro/halobgon.py +488 -0
  215. setiastro/saspro/header_viewer.py +448 -0
  216. setiastro/saspro/headless_utils.py +88 -0
  217. setiastro/saspro/histogram.py +756 -0
  218. setiastro/saspro/history_explorer.py +941 -0
  219. setiastro/saspro/i18n.py +168 -0
  220. setiastro/saspro/image_combine.py +417 -0
  221. setiastro/saspro/image_peeker_pro.py +1604 -0
  222. setiastro/saspro/imageops/__init__.py +37 -0
  223. setiastro/saspro/imageops/mdi_snap.py +292 -0
  224. setiastro/saspro/imageops/scnr.py +36 -0
  225. setiastro/saspro/imageops/starbasedwhitebalance.py +210 -0
  226. setiastro/saspro/imageops/stretch.py +236 -0
  227. setiastro/saspro/isophote.py +1182 -0
  228. setiastro/saspro/layers.py +208 -0
  229. setiastro/saspro/layers_dock.py +714 -0
  230. setiastro/saspro/lazy_imports.py +193 -0
  231. setiastro/saspro/legacy/__init__.py +2 -0
  232. setiastro/saspro/legacy/image_manager.py +2226 -0
  233. setiastro/saspro/legacy/numba_utils.py +3676 -0
  234. setiastro/saspro/legacy/xisf.py +1071 -0
  235. setiastro/saspro/linear_fit.py +537 -0
  236. setiastro/saspro/live_stacking.py +1841 -0
  237. setiastro/saspro/log_bus.py +5 -0
  238. setiastro/saspro/logging_config.py +460 -0
  239. setiastro/saspro/luminancerecombine.py +309 -0
  240. setiastro/saspro/main_helpers.py +201 -0
  241. setiastro/saspro/mask_creation.py +931 -0
  242. setiastro/saspro/masks_core.py +56 -0
  243. setiastro/saspro/mdi_widgets.py +353 -0
  244. setiastro/saspro/memory_utils.py +666 -0
  245. setiastro/saspro/metadata_patcher.py +75 -0
  246. setiastro/saspro/mfdeconv.py +3831 -0
  247. setiastro/saspro/mfdeconv_earlystop.py +71 -0
  248. setiastro/saspro/mfdeconvcudnn.py +3263 -0
  249. setiastro/saspro/mfdeconvsport.py +2382 -0
  250. setiastro/saspro/minorbodycatalog.py +567 -0
  251. setiastro/saspro/morphology.py +407 -0
  252. setiastro/saspro/multiscale_decomp.py +1293 -0
  253. setiastro/saspro/nbtorgb_stars.py +541 -0
  254. setiastro/saspro/numba_utils.py +3145 -0
  255. setiastro/saspro/numba_warmup.py +141 -0
  256. setiastro/saspro/ops/__init__.py +9 -0
  257. setiastro/saspro/ops/command_help_dialog.py +623 -0
  258. setiastro/saspro/ops/command_runner.py +217 -0
  259. setiastro/saspro/ops/commands.py +1594 -0
  260. setiastro/saspro/ops/script_editor.py +1102 -0
  261. setiastro/saspro/ops/scripts.py +1473 -0
  262. setiastro/saspro/ops/settings.py +637 -0
  263. setiastro/saspro/parallel_utils.py +554 -0
  264. setiastro/saspro/pedestal.py +121 -0
  265. setiastro/saspro/perfect_palette_picker.py +1071 -0
  266. setiastro/saspro/pipeline.py +110 -0
  267. setiastro/saspro/pixelmath.py +1604 -0
  268. setiastro/saspro/plate_solver.py +2445 -0
  269. setiastro/saspro/project_io.py +797 -0
  270. setiastro/saspro/psf_utils.py +136 -0
  271. setiastro/saspro/psf_viewer.py +549 -0
  272. setiastro/saspro/pyi_rthook_astroquery.py +95 -0
  273. setiastro/saspro/remove_green.py +331 -0
  274. setiastro/saspro/remove_stars.py +1599 -0
  275. setiastro/saspro/remove_stars_preset.py +404 -0
  276. setiastro/saspro/resources.py +501 -0
  277. setiastro/saspro/rgb_combination.py +208 -0
  278. setiastro/saspro/rgb_extract.py +19 -0
  279. setiastro/saspro/rgbalign.py +723 -0
  280. setiastro/saspro/runtime_imports.py +7 -0
  281. setiastro/saspro/runtime_torch.py +754 -0
  282. setiastro/saspro/save_options.py +73 -0
  283. setiastro/saspro/selective_color.py +1552 -0
  284. setiastro/saspro/sfcc.py +1472 -0
  285. setiastro/saspro/shortcuts.py +3043 -0
  286. setiastro/saspro/signature_insert.py +1102 -0
  287. setiastro/saspro/stacking_suite.py +18470 -0
  288. setiastro/saspro/star_alignment.py +7435 -0
  289. setiastro/saspro/star_alignment_preset.py +329 -0
  290. setiastro/saspro/star_metrics.py +49 -0
  291. setiastro/saspro/star_spikes.py +765 -0
  292. setiastro/saspro/star_stretch.py +507 -0
  293. setiastro/saspro/stat_stretch.py +538 -0
  294. setiastro/saspro/status_log_dock.py +78 -0
  295. setiastro/saspro/subwindow.py +3328 -0
  296. setiastro/saspro/supernovaasteroidhunter.py +1719 -0
  297. setiastro/saspro/swap_manager.py +99 -0
  298. setiastro/saspro/torch_backend.py +89 -0
  299. setiastro/saspro/torch_rejection.py +434 -0
  300. setiastro/saspro/translations/all_source_strings.json +3654 -0
  301. setiastro/saspro/translations/ar_translations.py +3865 -0
  302. setiastro/saspro/translations/de_translations.py +3749 -0
  303. setiastro/saspro/translations/es_translations.py +3939 -0
  304. setiastro/saspro/translations/fr_translations.py +3858 -0
  305. setiastro/saspro/translations/hi_translations.py +3571 -0
  306. setiastro/saspro/translations/integrate_translations.py +270 -0
  307. setiastro/saspro/translations/it_translations.py +3678 -0
  308. setiastro/saspro/translations/ja_translations.py +3601 -0
  309. setiastro/saspro/translations/pt_translations.py +3869 -0
  310. setiastro/saspro/translations/ru_translations.py +2848 -0
  311. setiastro/saspro/translations/saspro_ar.qm +0 -0
  312. setiastro/saspro/translations/saspro_ar.ts +255 -0
  313. setiastro/saspro/translations/saspro_de.qm +0 -0
  314. setiastro/saspro/translations/saspro_de.ts +253 -0
  315. setiastro/saspro/translations/saspro_es.qm +0 -0
  316. setiastro/saspro/translations/saspro_es.ts +12520 -0
  317. setiastro/saspro/translations/saspro_fr.qm +0 -0
  318. setiastro/saspro/translations/saspro_fr.ts +12514 -0
  319. setiastro/saspro/translations/saspro_hi.qm +0 -0
  320. setiastro/saspro/translations/saspro_hi.ts +257 -0
  321. setiastro/saspro/translations/saspro_it.qm +0 -0
  322. setiastro/saspro/translations/saspro_it.ts +12520 -0
  323. setiastro/saspro/translations/saspro_ja.qm +0 -0
  324. setiastro/saspro/translations/saspro_ja.ts +257 -0
  325. setiastro/saspro/translations/saspro_pt.qm +0 -0
  326. setiastro/saspro/translations/saspro_pt.ts +257 -0
  327. setiastro/saspro/translations/saspro_ru.qm +0 -0
  328. setiastro/saspro/translations/saspro_ru.ts +237 -0
  329. setiastro/saspro/translations/saspro_sw.qm +0 -0
  330. setiastro/saspro/translations/saspro_sw.ts +257 -0
  331. setiastro/saspro/translations/saspro_uk.qm +0 -0
  332. setiastro/saspro/translations/saspro_uk.ts +10771 -0
  333. setiastro/saspro/translations/saspro_zh.qm +0 -0
  334. setiastro/saspro/translations/saspro_zh.ts +12520 -0
  335. setiastro/saspro/translations/sw_translations.py +3671 -0
  336. setiastro/saspro/translations/uk_translations.py +3700 -0
  337. setiastro/saspro/translations/zh_translations.py +3675 -0
  338. setiastro/saspro/versioning.py +77 -0
  339. setiastro/saspro/view_bundle.py +1558 -0
  340. setiastro/saspro/wavescale_hdr.py +645 -0
  341. setiastro/saspro/wavescale_hdr_preset.py +101 -0
  342. setiastro/saspro/wavescalede.py +680 -0
  343. setiastro/saspro/wavescalede_preset.py +230 -0
  344. setiastro/saspro/wcs_update.py +374 -0
  345. setiastro/saspro/whitebalance.py +492 -0
  346. setiastro/saspro/widgets/__init__.py +48 -0
  347. setiastro/saspro/widgets/common_utilities.py +306 -0
  348. setiastro/saspro/widgets/graphics_views.py +122 -0
  349. setiastro/saspro/widgets/image_utils.py +518 -0
  350. setiastro/saspro/widgets/minigame/game.js +986 -0
  351. setiastro/saspro/widgets/minigame/index.html +53 -0
  352. setiastro/saspro/widgets/minigame/style.css +241 -0
  353. setiastro/saspro/widgets/preview_dialogs.py +280 -0
  354. setiastro/saspro/widgets/resource_monitor.py +237 -0
  355. setiastro/saspro/widgets/spinboxes.py +275 -0
  356. setiastro/saspro/widgets/themed_buttons.py +13 -0
  357. setiastro/saspro/widgets/wavelet_utils.py +331 -0
  358. setiastro/saspro/wimi.py +7996 -0
  359. setiastro/saspro/wims.py +578 -0
  360. setiastro/saspro/window_shelf.py +185 -0
  361. setiastro/saspro/xisf.py +1123 -0
  362. setiastrosuitepro-1.6.2.post1.dist-info/METADATA +278 -0
  363. setiastrosuitepro-1.6.2.post1.dist-info/RECORD +367 -0
  364. setiastrosuitepro-1.6.2.post1.dist-info/WHEEL +4 -0
  365. setiastrosuitepro-1.6.2.post1.dist-info/entry_points.txt +6 -0
  366. setiastrosuitepro-1.6.2.post1.dist-info/licenses/LICENSE +674 -0
  367. setiastrosuitepro-1.6.2.post1.dist-info/licenses/license.txt +2580 -0
@@ -0,0 +1,2382 @@
1
+ # pro/mfdeconvsport.py
2
+ from __future__ import annotations
3
+ import os, sys
4
+ import math
5
+ import re
6
+ import numpy as np
7
+ from astropy.io import fits
8
+ from PyQt6.QtCore import QObject, pyqtSignal
9
+ from setiastro.saspro.psf_utils import compute_psf_kernel_for_image
10
+ from PyQt6.QtWidgets import QApplication
11
+ from PyQt6.QtCore import QThread
12
+ from threadpoolctl import threadpool_limits
13
+ from concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor
14
+ _USE_PROCESS_POOL_FOR_ASSETS = not getattr(sys, "frozen", False)
15
+ from setiastro.saspro.mfdeconv_earlystop import EarlyStopper
16
+
17
+ import contextlib
18
+ try:
19
+ import sep
20
+ except Exception:
21
+ sep = None
22
+ from setiastro.saspro.free_torch_memory import _free_torch_memory
23
+ torch = None # filled by runtime loader if available
24
+ TORCH_OK = False
25
+ NO_GRAD = contextlib.nullcontext # fallback
26
+
27
+ _XISF_READERS = []
28
+ try:
29
+ # e.g. your legacy module
30
+ from setiastro.saspro.legacy import xisf as _legacy_xisf
31
+ if hasattr(_legacy_xisf, "read"):
32
+ _XISF_READERS.append(lambda p: _legacy_xisf.read(p))
33
+ elif hasattr(_legacy_xisf, "open"):
34
+ _XISF_READERS.append(lambda p: _legacy_xisf.open(p)[0])
35
+ except Exception:
36
+ pass
37
+ try:
38
+ # sometimes projects expose a generic load_image
39
+ from setiastro.saspro.legacy.image_manager import load_image as _generic_load_image # adjust if needed
40
+ _XISF_READERS.append(lambda p: _generic_load_image(p)[0])
41
+ except Exception:
42
+ pass
43
+
44
+ # at top of file with the other imports
45
+ from concurrent.futures import ThreadPoolExecutor, as_completed
46
+ from queue import SimpleQueue
47
+ from setiastro.saspro.memory_utils import LRUDict
48
+
49
+ # ── XISF decode cache → memmap on disk ─────────────────────────────────
50
+ import tempfile
51
+ import threading
52
+ import uuid
53
+ import atexit
54
+ _XISF_CACHE = LRUDict(50)
55
+ _XISF_LOCK = threading.Lock()
56
+ _XISF_TMPFILES = []
57
+
58
+ from collections import OrderedDict
59
+
60
+ # ─────────────────────────────────────────────────────────────────────────────
61
+ # Unified image I/O for MFDeconv (FITS + XISF)
62
+ # ─────────────────────────────────────────────────────────────────────────────
63
+ import os
64
+ import numpy as np
65
+ from astropy.io import fits
66
+
67
+ from pathlib import Path
68
+
69
+
70
+ from collections import OrderedDict
71
+
72
+ # ── CHW LRU (float32) built on top of FITS memmap & XISF memmap ────────────────
73
+ class _FrameCHWLRU:
74
+ def __init__(self, capacity=8):
75
+ self.cap = int(max(1, capacity))
76
+ self.od = OrderedDict()
77
+
78
+ def clear(self):
79
+ self.od.clear()
80
+
81
+ def get(self, path, Ht, Wt, color_mode):
82
+ key = (path, Ht, Wt, str(color_mode).lower())
83
+ hit = self.od.get(key)
84
+ if hit is not None:
85
+ self.od.move_to_end(key)
86
+ return hit
87
+
88
+ # Load backing array cheaply (memmap for FITS, cached memmap for XISF)
89
+ ext = os.path.splitext(path)[1].lower()
90
+ if ext == ".xisf":
91
+ a = _xisf_cached_array(path) # float32, HW/HWC/CHW
92
+ else:
93
+ # FITS path: use astropy memmap (no data copy)
94
+ with fits.open(path, memmap=True, ignore_missing_simple=True) as hdul:
95
+ arr = None
96
+ for h in hdul:
97
+ if getattr(h, "data", None) is not None:
98
+ arr = h.data
99
+ break
100
+ if arr is None:
101
+ raise ValueError(f"No image data in {path}")
102
+ a = np.asarray(arr)
103
+ # dtype normalize once; keep float32
104
+ if a.dtype.kind in "ui":
105
+ a = a.astype(np.float32) / (float(np.iinfo(a.dtype).max) or 1.0)
106
+ else:
107
+ a = a.astype(np.float32, copy=False)
108
+
109
+ # Center-crop to (Ht, Wt) and convert to CHW
110
+ a = np.asarray(a) # float32
111
+ a = _center_crop(a, Ht, Wt)
112
+
113
+ # Respect color_mode: “luma” → 1×H×W, “PerChannel” → 3×H×W if RGB present
114
+ cm = str(color_mode).lower()
115
+ if cm == "luma":
116
+ a_chw = _as_chw(_to_luma_local(a)).astype(np.float32, copy=False)
117
+ else:
118
+ a_chw = _as_chw(a).astype(np.float32, copy=False)
119
+ if a_chw.shape[0] == 1 and cm != "luma":
120
+ # still OK (mono data)
121
+ pass
122
+
123
+ # LRU insert
124
+ self.od[key] = a_chw
125
+ if len(self.od) > self.cap:
126
+ self.od.popitem(last=False)
127
+ return a_chw
128
+
129
+ _FRAME_LRU = _FrameCHWLRU(capacity=8) # tune if you like
130
+
131
+ def _clear_all_caches():
132
+ try: _clear_xisf_cache()
133
+ except Exception as e:
134
+ import logging
135
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
136
+ try: _FRAME_LRU.clear()
137
+ except Exception as e:
138
+ import logging
139
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
140
+
141
+ def _as_chw(np_img: np.ndarray) -> np.ndarray:
142
+ x = np.asarray(np_img, dtype=np.float32, order="C")
143
+ if x.size == 0:
144
+ raise RuntimeError(f"Empty image array after load; raw shape={np_img.shape}")
145
+ if x.ndim == 2:
146
+ return x[None, ...] # 1,H,W
147
+ if x.ndim == 3 and x.shape[0] in (1, 3):
148
+ if x.shape[0] == 0:
149
+ raise RuntimeError(f"Zero channels in CHW array; shape={x.shape}")
150
+ return x
151
+ if x.ndim == 3 and x.shape[-1] in (1, 3):
152
+ if x.shape[-1] == 0:
153
+ raise RuntimeError(f"Zero channels in HWC array; shape={x.shape}")
154
+ return np.moveaxis(x, -1, 0)
155
+ # last resort: treat first dim as channels, but reject zero
156
+ if x.shape[0] == 0:
157
+ raise RuntimeError(f"Zero channels in array; shape={x.shape}")
158
+ return x
159
+
160
+ def _normalize_to_float32(a: np.ndarray) -> np.ndarray:
161
+ if a.dtype.kind in "ui":
162
+ return (a.astype(np.float32) / (float(np.iinfo(a.dtype).max) or 1.0))
163
+ if a.dtype == np.float32:
164
+ return a
165
+ return a.astype(np.float32, copy=False)
166
+
167
+ def _xisf_cached_array(path: str) -> np.memmap:
168
+ """
169
+ Decode an XISF image exactly once and back it by a read-only float32 memmap.
170
+ Returns a memmap that can be sliced cheaply for tiles.
171
+ """
172
+ with _XISF_LOCK:
173
+ hit = _XISF_CACHE.get(path)
174
+ if hit is not None:
175
+ fn, shape = hit
176
+ return np.memmap(fn, dtype=np.float32, mode="r", shape=shape)
177
+
178
+ # Decode once
179
+ arr, _ = _load_image_array(path) # your existing loader
180
+ if arr is None:
181
+ raise ValueError(f"XISF loader returned None for {path}")
182
+ arr = np.asarray(arr)
183
+ arrf = _normalize_to_float32(arr)
184
+
185
+ # Create a temp file-backed memmap
186
+ tmpdir = tempfile.gettempdir()
187
+ fn = os.path.join(tmpdir, f"xisf_cache_{uuid.uuid4().hex}.mmap")
188
+ mm = np.memmap(fn, dtype=np.float32, mode="w+", shape=arrf.shape)
189
+ mm[...] = arrf[...]
190
+ mm.flush()
191
+ del mm # close writer handle; re-open below as read-only
192
+
193
+ _XISF_CACHE[path] = (fn, arrf.shape)
194
+ _XISF_TMPFILES.append(fn)
195
+ return np.memmap(fn, dtype=np.float32, mode="r", shape=arrf.shape)
196
+
197
+ def _clear_xisf_cache():
198
+ with _XISF_LOCK:
199
+ for fn in _XISF_TMPFILES:
200
+ try: os.remove(fn)
201
+ except Exception as e:
202
+ import logging
203
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
204
+ _XISF_CACHE.clear()
205
+ _XISF_TMPFILES.clear()
206
+
207
+ atexit.register(_clear_xisf_cache)
208
+
209
+
210
+ def _is_xisf(path: str) -> bool:
211
+ return os.path.splitext(path)[1].lower() == ".xisf"
212
+
213
+ def _read_xisf_numpy(path: str) -> np.ndarray:
214
+ if not _XISF_READERS:
215
+ raise RuntimeError(
216
+ "No XISF readers registered. Ensure one of "
217
+ "legacy.xisf.read/open or *.image_io.load_image is importable."
218
+ )
219
+ last_err = None
220
+ for fn in _XISF_READERS:
221
+ try:
222
+ arr = fn(path)
223
+ if isinstance(arr, tuple):
224
+ arr = arr[0]
225
+ return np.asarray(arr)
226
+ except Exception as e:
227
+ last_err = e
228
+ raise RuntimeError(f"All XISF readers failed for {path}: {last_err}")
229
+
230
+ def _fits_open_data(path: str):
231
+ # ignore_missing_simple=True lets us open headers missing SIMPLE
232
+ with fits.open(path, memmap=True, ignore_missing_simple=True) as hdul:
233
+ hdu = hdul[0]
234
+ if hdu.data is None:
235
+ # find first image HDU if primary is header-only
236
+ for h in hdul[1:]:
237
+ if getattr(h, "data", None) is not None:
238
+ hdu = h
239
+ break
240
+ data = np.asanyarray(hdu.data)
241
+ hdr = hdu.header
242
+ return data, hdr
243
+
244
+ def _load_image_array(path: str) -> tuple[np.ndarray, "fits.Header | None"]:
245
+ """
246
+ Return (numpy array, fits.Header or None). Color-last if 3D.
247
+ dtype left as-is; callers cast to float32. Array is C-contig & writeable.
248
+ """
249
+ if _is_xisf(path):
250
+ arr = _read_xisf_numpy(path)
251
+ hdr = None
252
+ else:
253
+ arr, hdr = _fits_open_data(path)
254
+
255
+ a = np.asarray(arr)
256
+ # Move color axis to last if 3D with a leading channel axis
257
+ if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3):
258
+ a = np.moveaxis(a, 0, -1)
259
+ # Ensure contiguous, writeable float32 decisions happen later; here we just ensure writeable
260
+ if (not a.flags.c_contiguous) or (not a.flags.writeable):
261
+ a = np.array(a, copy=True)
262
+ return a, hdr
263
+
264
+ def _probe_hw(path: str) -> tuple[int, int, int | None]:
265
+ """
266
+ Returns (H, W, C_or_None) without changing data. Moves color to last if needed.
267
+ """
268
+ a, _ = _load_image_array(path)
269
+ if a.ndim == 2:
270
+ return a.shape[0], a.shape[1], None
271
+ if a.ndim == 3:
272
+ h, w, c = a.shape
273
+ # treat mono-3D as (H,W,1)
274
+ if c not in (1, 3) and a.shape[0] in (1, 3):
275
+ a = np.moveaxis(a, 0, -1)
276
+ h, w, c = a.shape
277
+ return h, w, c if c in (1, 3) else None
278
+ raise ValueError(f"Unsupported ndim={a.ndim} for {path}")
279
+
280
+ def _common_hw_from_paths(paths: list[str]) -> tuple[int, int]:
281
+ """
282
+ Replacement for the old FITS-only version: min(H), min(W) across files.
283
+ """
284
+ Hs, Ws = [], []
285
+ for p in paths:
286
+ h, w, _ = _probe_hw(p)
287
+ Hs.append(int(h)); Ws.append(int(w))
288
+ return int(min(Hs)), int(min(Ws))
289
+
290
+ def _to_chw_float32(img: np.ndarray, color_mode: str) -> np.ndarray:
291
+ """
292
+ Convert to CHW float32:
293
+ - mono → (1,H,W)
294
+ - RGB → (3,H,W) if 'PerChannel'; (1,H,W) if 'luma'
295
+ """
296
+ x = np.asarray(img)
297
+ if x.ndim == 2:
298
+ y = x.astype(np.float32, copy=False)[None, ...] # (1,H,W)
299
+ return y
300
+ if x.ndim == 3:
301
+ # color-last (H,W,C) expected
302
+ if x.shape[-1] == 1:
303
+ return x[..., 0].astype(np.float32, copy=False)[None, ...]
304
+ if x.shape[-1] == 3:
305
+ if str(color_mode).lower() in ("perchannel", "per_channel", "perchannelrgb"):
306
+ r, g, b = x[..., 0], x[..., 1], x[..., 2]
307
+ return np.stack([r.astype(np.float32, copy=False),
308
+ g.astype(np.float32, copy=False),
309
+ b.astype(np.float32, copy=False)], axis=0)
310
+ # luma
311
+ r, g, b = x[..., 0].astype(np.float32, copy=False), x[..., 1].astype(np.float32, copy=False), x[..., 2].astype(np.float32, copy=False)
312
+ L = 0.2126*r + 0.7152*g + 0.0722*b
313
+ return L[None, ...]
314
+ # rare mono-3D
315
+ if x.shape[0] in (1, 3) and x.shape[-1] not in (1, 3):
316
+ x = np.moveaxis(x, 0, -1)
317
+ return _to_chw_float32(x, color_mode)
318
+ raise ValueError(f"Unsupported image shape {x.shape}")
319
+
320
+ def _center_crop_hw(img: np.ndarray, Ht: int, Wt: int) -> np.ndarray:
321
+ h, w = img.shape[:2]
322
+ y0 = max(0, (h - Ht)//2); x0 = max(0, (w - Wt)//2)
323
+ return img[y0:y0+Ht, x0:x0+Wt, ...].copy() if (Ht < h or Wt < w) else img
324
+
325
+ def _stack_loader_memmap(paths: list[str], Ht: int, Wt: int, color_mode: str):
326
+ """
327
+ Drop-in replacement of the old FITS-only helper.
328
+ Returns (ys, hdrs):
329
+ ys : list of CHW float32 arrays cropped to (Ht,Wt)
330
+ hdrs : list of fits.Header or None (XISF)
331
+ """
332
+ ys, hdrs = [], []
333
+ for p in paths:
334
+ arr, hdr = _load_image_array(p)
335
+ arr = _center_crop_hw(arr, Ht, Wt)
336
+ # normalize integer data to [0,1] like the rest of your code
337
+ if arr.dtype.kind in "ui":
338
+ mx = np.float32(np.iinfo(arr.dtype).max)
339
+ arr = arr.astype(np.float32, copy=False) / (mx if mx > 0 else 1.0)
340
+ elif arr.dtype.kind == "f":
341
+ arr = arr.astype(np.float32, copy=False)
342
+ else:
343
+ arr = arr.astype(np.float32, copy=False)
344
+
345
+ y = _to_chw_float32(arr, color_mode)
346
+ if (not y.flags.c_contiguous) or (not y.flags.writeable):
347
+ y = np.ascontiguousarray(y.astype(np.float32, copy=True))
348
+ ys.append(y)
349
+ hdrs.append(hdr if isinstance(hdr, fits.Header) else None)
350
+ return ys, hdrs
351
+
352
+ def _safe_primary_header(path: str) -> fits.Header:
353
+ if _is_xisf(path):
354
+ # best-effort synthetic header
355
+ h = fits.Header()
356
+ h["SIMPLE"] = (True, "created by MFDeconv")
357
+ h["BITPIX"] = -32
358
+ h["NAXIS"] = 2
359
+ return h
360
+ try:
361
+ return fits.getheader(path, ext=0, ignore_missing_simple=True)
362
+ except Exception:
363
+ return fits.Header()
364
+
365
+
366
+ def _compute_frame_assets(i, arr, hdr, *, make_masks, make_varmaps,
367
+ star_mask_cfg, varmap_cfg, status_sink=lambda s: None):
368
+ """
369
+ Worker function: compute PSF and optional star mask / varmap for one frame.
370
+ Returns (index, psf, mask_or_None, var_or_None, log_lines)
371
+ """
372
+ logs = []
373
+ def log(s): logs.append(s)
374
+
375
+ # --- PSF sizing by FWHM ---
376
+ f_hdr = _estimate_fwhm_from_header(hdr)
377
+ f_img = _estimate_fwhm_from_image(arr)
378
+ f_whm = f_hdr if (np.isfinite(f_hdr)) else f_img
379
+ if not np.isfinite(f_whm) or f_whm <= 0:
380
+ f_whm = 2.5
381
+ k_auto = _auto_ksize_from_fwhm(f_whm)
382
+
383
+ # --- Star-derived PSF with retries ---
384
+ tried, psf = [], None
385
+ for k_try in [k_auto, max(k_auto - 4, 11), 21, 17, 15, 13, 11]:
386
+ if k_try in tried: continue
387
+ tried.append(k_try)
388
+ try:
389
+ out = compute_psf_kernel_for_image(arr, ksize=k_try, det_sigma=6.0, max_stars=80)
390
+ psf_try = out[0] if (isinstance(out, tuple) and len(out) >= 1) else out
391
+ if psf_try is not None:
392
+ psf = psf_try
393
+ break
394
+ except Exception:
395
+ psf = None
396
+ if psf is None:
397
+ psf = _gaussian_psf(f_whm, ksize=k_auto)
398
+ psf = _soften_psf(_normalize_psf(psf.astype(np.float32, copy=False)), sigma_px=0.25)
399
+
400
+ mask = None
401
+ var = None
402
+
403
+ if make_masks or make_varmaps:
404
+ # one background per frame (reused by both)
405
+ luma = _to_luma_local(arr)
406
+ vmc = (varmap_cfg or {})
407
+ sky_map, rms_map, err_scalar = _sep_background_precompute(
408
+ luma, bw=int(vmc.get("bw", 64)), bh=int(vmc.get("bh", 64))
409
+ )
410
+
411
+ if make_masks:
412
+ smc = star_mask_cfg or {}
413
+ mask = _star_mask_from_precomputed(
414
+ luma, sky_map, err_scalar,
415
+ thresh_sigma = smc.get("thresh_sigma", THRESHOLD_SIGMA),
416
+ max_objs = smc.get("max_objs", STAR_MASK_MAXOBJS),
417
+ grow_px = smc.get("grow_px", GROW_PX),
418
+ ellipse_scale= smc.get("ellipse_scale", ELLIPSE_SCALE),
419
+ soft_sigma = smc.get("soft_sigma", SOFT_SIGMA),
420
+ max_radius_px= smc.get("max_radius_px", MAX_STAR_RADIUS),
421
+ keep_floor = smc.get("keep_floor", KEEP_FLOOR),
422
+ max_side = smc.get("max_side", STAR_MASK_MAXSIDE),
423
+ status_cb = log,
424
+ )
425
+
426
+ if make_varmaps:
427
+ vmc = varmap_cfg or {}
428
+ var = _variance_map_from_precomputed(
429
+ luma, sky_map, rms_map, hdr,
430
+ smooth_sigma = vmc.get("smooth_sigma", 1.0),
431
+ floor = vmc.get("floor", 1e-8),
432
+ status_cb = log,
433
+ )
434
+
435
+ # small per-frame summary
436
+ fwhm_est = _psf_fwhm_px(psf)
437
+ logs.insert(0, f"MFDeconv: PSF{i}: ksize={psf.shape[0]} | FWHM≈{fwhm_est:.2f}px")
438
+
439
+ return i, psf, mask, var, logs
440
+
441
+ def _compute_one_worker(args):
442
+ """
443
+ Top-level picklable worker for ProcessPoolExecutor.
444
+ args: (i, path, make_masks_in_worker, make_varmaps, star_mask_cfg, varmap_cfg)
445
+ Returns (i, psf, mask, var, logs)
446
+ """
447
+ (i, path, make_masks_in_worker, make_varmaps, star_mask_cfg, varmap_cfg) = args
448
+ # avoid BLAS/OMP storm inside each process
449
+ with threadpool_limits(limits=1):
450
+ arr, hdr = _load_image_array(path) # FITS or XISF
451
+ arr = np.asarray(arr, dtype=np.float32, order="C")
452
+ if arr.ndim == 3 and arr.shape[-1] == 1:
453
+ arr = np.squeeze(arr, axis=-1)
454
+ if not isinstance(hdr, fits.Header): # synthesize FITS-like header for XISF
455
+ hdr = _safe_primary_header(path)
456
+ return _compute_frame_assets(
457
+ i, arr, hdr,
458
+ make_masks=bool(make_masks_in_worker),
459
+ make_varmaps=bool(make_varmaps),
460
+ star_mask_cfg=star_mask_cfg,
461
+ varmap_cfg=varmap_cfg,
462
+ )
463
+
464
+
465
+ def _build_psf_and_assets(
466
+ paths, # list[str]
467
+ make_masks=False,
468
+ make_varmaps=False,
469
+ status_cb=lambda s: None,
470
+ save_dir: str | None = None,
471
+ star_mask_cfg: dict | None = None,
472
+ varmap_cfg: dict | None = None,
473
+ max_workers: int | None = None,
474
+ star_mask_ref_path: str | None = None, # build one mask from this frame if provided
475
+ # NEW (passed from multiframe_deconv so we don’t re-probe/convert):
476
+ Ht: int | None = None,
477
+ Wt: int | None = None,
478
+ color_mode: str = "luma",
479
+ ):
480
+ """
481
+ Parallel PSF + (optional) star mask + variance map per frame.
482
+
483
+ Changes from the original:
484
+ • Reuses the decoded frame cache (_FRAME_LRU) for FITS/XISF so we never re-decode.
485
+ • Automatically switches to threads for XISF (so memmaps are shared across workers).
486
+ • Builds a single reference star mask (if requested) from the cached frame and
487
+ center-pads/crops it for all frames (no extra I/O).
488
+ • Preserves return order and streams worker logs back to the UI.
489
+ """
490
+ if save_dir:
491
+ os.makedirs(save_dir, exist_ok=True)
492
+
493
+ n = len(paths)
494
+
495
+ # Resolve target intersection size if caller didn't pass it
496
+ if Ht is None or Wt is None:
497
+ Ht, Wt = _common_hw_from_paths(paths)
498
+
499
+ # Sensible default worker count (cap at 8)
500
+ if max_workers is None:
501
+ try:
502
+ hw = os.cpu_count() or 4
503
+ except Exception:
504
+ hw = 4
505
+ max_workers = max(1, min(8, hw))
506
+
507
+ # Decide executor: for any XISF, prefer threads so the memmap/cache is shared
508
+ any_xisf = any(os.path.splitext(p)[1].lower() == ".xisf" for p in paths)
509
+ use_proc_pool = (not any_xisf) and _USE_PROCESS_POOL_FOR_ASSETS
510
+ Executor = ProcessPoolExecutor if use_proc_pool else ThreadPoolExecutor
511
+ pool_kind = "process" if use_proc_pool else "thread"
512
+ status_cb(f"MFDeconv: measuring PSFs/masks/varmaps with {max_workers} {pool_kind}s…")
513
+
514
+ # ---- helper: pad-or-crop a 2D array to (Ht,Wt), centered ----
515
+ def _center_pad_or_crop_2d(a2d: np.ndarray, Ht: int, Wt: int, fill: float = 1.0) -> np.ndarray:
516
+ a2d = np.asarray(a2d, dtype=np.float32)
517
+ H, W = int(a2d.shape[0]), int(a2d.shape[1])
518
+ # crop first if bigger
519
+ y0 = max(0, (H - Ht) // 2); x0 = max(0, (W - Wt) // 2)
520
+ y1 = min(H, y0 + Ht); x1 = min(W, x0 + Wt)
521
+ cropped = a2d[y0:y1, x0:x1]
522
+ ch, cw = cropped.shape
523
+ if ch == Ht and cw == Wt:
524
+ return np.ascontiguousarray(cropped, dtype=np.float32)
525
+ # pad if smaller
526
+ out = np.full((Ht, Wt), float(fill), dtype=np.float32)
527
+ oy = (Ht - ch) // 2; ox = (Wt - cw) // 2
528
+ out[oy:oy+ch, ox:ox+cw] = cropped
529
+ return out
530
+
531
+ # ---- optional: build one mask from the reference frame and reuse ----
532
+ base_ref_mask = None
533
+ if make_masks and star_mask_ref_path:
534
+ try:
535
+ status_cb(f"Star mask: using reference frame for all masks → {os.path.basename(star_mask_ref_path)}")
536
+ # Pull from the shared frame cache as luma on (Ht,Wt)
537
+ ref_chw = _FRAME_LRU.get(star_mask_ref_path, Ht, Wt, "luma") # (1,H,W) or (H,W)
538
+ L = ref_chw[0] if (ref_chw.ndim == 3) else ref_chw # 2D float32
539
+
540
+ vmc = (varmap_cfg or {})
541
+ sky_map, rms_map, err_scalar = _sep_background_precompute(
542
+ L, bw=int(vmc.get("bw", 64)), bh=int(vmc.get("bh", 64))
543
+ )
544
+ smc = (star_mask_cfg or {})
545
+ base_ref_mask = _star_mask_from_precomputed(
546
+ L, sky_map, err_scalar,
547
+ thresh_sigma = smc.get("thresh_sigma", THRESHOLD_SIGMA),
548
+ max_objs = smc.get("max_objs", STAR_MASK_MAXOBJS),
549
+ grow_px = smc.get("grow_px", GROW_PX),
550
+ ellipse_scale= smc.get("ellipse_scale", ELLIPSE_SCALE),
551
+ soft_sigma = smc.get("soft_sigma", SOFT_SIGMA),
552
+ max_radius_px= smc.get("max_radius_px", MAX_STAR_RADIUS),
553
+ keep_floor = smc.get("keep_floor", KEEP_FLOOR),
554
+ max_side = smc.get("max_side", STAR_MASK_MAXSIDE),
555
+ status_cb = status_cb,
556
+ )
557
+ except Exception as e:
558
+ status_cb(f"⚠️ Star mask (reference) failed: {e}. Falling back to per-frame masks.")
559
+ base_ref_mask = None
560
+
561
+ # for GUI safety, queue logs from workers and flush in the main thread
562
+ log_queue: SimpleQueue = SimpleQueue()
563
+
564
+ def enqueue_logs(lines):
565
+ for s in lines:
566
+ log_queue.put(s)
567
+
568
+ psfs = [None] * n
569
+ masks = ([None] * n) if make_masks else None
570
+ vars_ = ([None] * n) if make_varmaps else None
571
+ make_masks_in_worker = bool(make_masks and (base_ref_mask is None))
572
+
573
+ # --- thread worker: get frame from cache and compute assets ---
574
+ def _compute_one(i: int, path: str):
575
+ # avoid heavy BLAS oversubscription inside each worker
576
+ with threadpool_limits(limits=1):
577
+ # Pull frame from cache honoring color_mode & target (Ht,Wt)
578
+ img_chw = _FRAME_LRU.get(path, Ht, Wt, color_mode) # (C,H,W) float32
579
+ # For PSF/mask/varmap we operate on a 2D plane (luma/mono)
580
+ arr2d = img_chw[0] if (img_chw.ndim == 3) else img_chw # (H,W) float32
581
+
582
+ # Header: synthesize a safe FITS-like header (works for XISF too)
583
+ try:
584
+ hdr = _safe_primary_header(path)
585
+ except Exception:
586
+ hdr = fits.Header()
587
+
588
+ return _compute_frame_assets(
589
+ i, arr2d, hdr,
590
+ make_masks=bool(make_masks_in_worker),
591
+ make_varmaps=bool(make_varmaps),
592
+ star_mask_cfg=star_mask_cfg,
593
+ varmap_cfg=varmap_cfg,
594
+ )
595
+
596
+ # --- submit jobs ---
597
+ with Executor(max_workers=max_workers) as ex:
598
+ futs = []
599
+ for i, p in enumerate(paths, start=1):
600
+ status_cb(f"MFDeconv: measuring PSF {i}/{n} …")
601
+ if use_proc_pool:
602
+ # Process-safe path: worker re-loads inside the subprocess
603
+ futs.append(ex.submit(
604
+ _compute_one_worker,
605
+ (i, p, bool(make_masks_in_worker), bool(make_varmaps), star_mask_cfg, varmap_cfg)
606
+ ))
607
+ else:
608
+ # Thread path: hits the shared cache (fast path for XISF/FITS)
609
+ futs.append(ex.submit(_compute_one, i, p))
610
+
611
+ done_cnt = 0
612
+ for fut in as_completed(futs):
613
+ i, psf, m, v, logs = fut.result()
614
+ idx = i - 1
615
+ psfs[idx] = psf
616
+ if masks is not None:
617
+ masks[idx] = m
618
+ if vars_ is not None:
619
+ vars_[idx] = v
620
+ enqueue_logs(logs)
621
+
622
+ done_cnt += 1
623
+ if (done_cnt % 4) == 0 or done_cnt == n:
624
+ while not log_queue.empty():
625
+ try:
626
+ status_cb(log_queue.get_nowait())
627
+ except Exception:
628
+ break
629
+
630
+ # If we built a single reference mask, apply it to every frame (center pad/crop)
631
+ if base_ref_mask is not None and masks is not None:
632
+ for idx in range(n):
633
+ masks[idx] = _center_pad_or_crop_2d(base_ref_mask, int(Ht), int(Wt), fill=1.0)
634
+
635
+ # final flush of any remaining logs
636
+ while not log_queue.empty():
637
+ try:
638
+ status_cb(log_queue.get_nowait())
639
+ except Exception:
640
+ break
641
+
642
+ # save PSFs if requested
643
+ if save_dir:
644
+ for i, k in enumerate(psfs, start=1):
645
+ if k is not None:
646
+ fits.PrimaryHDU(k.astype(np.float32, copy=False)).writeto(
647
+ os.path.join(save_dir, f"psf_{i:03d}.fit"), overwrite=True
648
+ )
649
+
650
+ return psfs, masks, vars_
651
+
652
+ _ALLOWED = re.compile(r"[^A-Za-z0-9_-]+")
653
+
654
+ # known FITS-style multi-extensions (rightmost-first match)
655
+ _KNOWN_EXTS = [
656
+ ".fits.fz", ".fit.fz", ".fits.gz", ".fit.gz",
657
+ ".fz", ".gz",
658
+ ".fits", ".fit"
659
+ ]
660
+
661
+ def _sanitize_token(s: str) -> str:
662
+ s = _ALLOWED.sub("_", s)
663
+ s = re.sub(r"_+", "_", s).strip("_")
664
+ return s
665
+
666
+ def _split_known_exts(p: Path) -> tuple[str, str]:
667
+ """
668
+ Return (name_body, full_ext) where full_ext is a REAL extension block
669
+ (e.g. '.fits.fz'). Any junk like '.0s (1310x880)_MFDeconv' stays in body.
670
+ """
671
+ name = p.name
672
+ for ext in _KNOWN_EXTS:
673
+ if name.lower().endswith(ext):
674
+ body = name[:-len(ext)]
675
+ return body, ext
676
+ # fallback: single suffix
677
+ return p.stem, "".join(p.suffixes)
678
+
679
+ _SIZE_RE = re.compile(r"\(?\s*(\d{2,5})x(\d{2,5})\s*\)?", re.IGNORECASE)
680
+ _EXP_RE = re.compile(r"(?<![A-Za-z0-9])(\d+(?:\.\d+)?)\s*s\b", re.IGNORECASE)
681
+ _RX_RE = re.compile(r"(?<![A-Za-z0-9])(\d+)x\b", re.IGNORECASE)
682
+
683
+ def _extract_size(body: str) -> str | None:
684
+ m = _SIZE_RE.search(body)
685
+ return f"{m.group(1)}x{m.group(2)}" if m else None
686
+
687
+ def _extract_exposure_secs(body: str) -> str | None:
688
+ m = _EXP_RE.search(body)
689
+ if not m:
690
+ return None
691
+ secs = int(round(float(m.group(1))))
692
+ return f"{secs}s"
693
+
694
+ def _strip_metadata_from_base(body: str) -> str:
695
+ s = body
696
+
697
+ # normalize common separators first
698
+ s = s.replace(" - ", "_")
699
+
700
+ # remove known trailing marker '_MFDeconv'
701
+ s = re.sub(r"(?i)[\s_]+MFDeconv$", "", s)
702
+
703
+ # remove parenthetical copy counters e.g. '(1)'
704
+ s = re.sub(r"\(\s*\d+\s*\)$", "", s)
705
+
706
+ # remove size (with or without parens) anywhere
707
+ s = _SIZE_RE.sub("", s)
708
+
709
+ # remove exposures like '0s', '0.5s', ' 45 s' (even if preceded by a dot)
710
+ s = _EXP_RE.sub("", s)
711
+
712
+ # remove any _#x tokens
713
+ s = _RX_RE.sub("", s)
714
+
715
+ # collapse whitespace/underscores and sanitize
716
+ s = re.sub(r"[\s]+", "_", s)
717
+ s = _sanitize_token(s)
718
+ return s or "output"
719
+
720
+ def _canonical_out_name_prefix(base: str, r: int, size: str | None,
721
+ exposure_secs: str | None, tag: str = "MFDeconv") -> str:
722
+ parts = [_sanitize_token(tag), _sanitize_token(base)]
723
+ if size:
724
+ parts.append(_sanitize_token(size))
725
+ if exposure_secs:
726
+ parts.append(_sanitize_token(exposure_secs))
727
+ if int(max(1, r)) > 1:
728
+ parts.append(f"{int(r)}x")
729
+ return "_".join(parts)
730
+
731
+ def _sr_out_path(out_path: str, r: int) -> Path:
732
+ """
733
+ Build: MFDeconv_<base>[_<HxW>][_<secs>s][_2x], preserving REAL extensions.
734
+ """
735
+ p = Path(out_path)
736
+ body, real_ext = _split_known_exts(p)
737
+
738
+ # harvest metadata from the whole body (not Path.stem)
739
+ size = _extract_size(body)
740
+ ex_sec = _extract_exposure_secs(body)
741
+
742
+ # clean base
743
+ base = _strip_metadata_from_base(body)
744
+
745
+ new_stem = _canonical_out_name_prefix(base, r=int(max(1, r)), size=size, exposure_secs=ex_sec, tag="MFDeconv")
746
+ return p.with_name(f"{new_stem}{real_ext}")
747
+
748
+ def _nonclobber_path(path: str) -> str:
749
+ """
750
+ Version collisions as '_v2', '_v3', ... (no spaces/parentheses).
751
+ """
752
+ p = Path(path)
753
+ if not p.exists():
754
+ return str(p)
755
+
756
+ # keep the true extension(s)
757
+ body, real_ext = _split_known_exts(p)
758
+
759
+ # if already has _vN, bump it
760
+ m = re.search(r"(.*)_v(\d+)$", body)
761
+ if m:
762
+ base = m.group(1); n = int(m.group(2)) + 1
763
+ else:
764
+ base = body; n = 2
765
+
766
+ while True:
767
+ candidate = p.with_name(f"{base}_v{n}{real_ext}")
768
+ if not candidate.exists():
769
+ return str(candidate)
770
+ n += 1
771
+
772
+ def _iter_folder(basefile: str) -> str:
773
+ d, fname = os.path.split(basefile)
774
+ root, ext = os.path.splitext(fname)
775
+ tgt = os.path.join(d, f"{root}.iters")
776
+ if not os.path.exists(tgt):
777
+ try:
778
+ os.makedirs(tgt, exist_ok=True)
779
+ except Exception:
780
+ # last resort: suffix (n)
781
+ n = 1
782
+ while True:
783
+ cand = os.path.join(d, f"{root}.iters ({n})")
784
+ try:
785
+ os.makedirs(cand, exist_ok=True)
786
+ return cand
787
+ except Exception:
788
+ n += 1
789
+ return tgt
790
+
791
+ def _save_iter_image(arr, hdr_base, folder, tag, color_mode):
792
+ """
793
+ arr: numpy array (H,W) or (C,H,W) float32
794
+ tag: 'seed' or 'iter_###'
795
+ """
796
+ if arr.ndim == 3 and arr.shape[0] not in (1, 3) and arr.shape[-1] in (1, 3):
797
+ arr = np.moveaxis(arr, -1, 0)
798
+ if arr.ndim == 3 and arr.shape[0] == 1:
799
+ arr = arr[0]
800
+
801
+ hdr = fits.Header(hdr_base) if isinstance(hdr_base, fits.Header) else fits.Header()
802
+ hdr['MF_PART'] = (str(tag), 'MFDeconv intermediate (seed/iter)')
803
+ hdr['MF_COLOR'] = (str(color_mode), 'Color mode used')
804
+ path = os.path.join(folder, f"{tag}.fit")
805
+ # overwrite allowed inside the dedicated folder
806
+ fits.PrimaryHDU(data=arr.astype(np.float32, copy=False), header=hdr).writeto(path, overwrite=True)
807
+ return path
808
+
809
+
810
+ def _process_gui_events_safely():
811
+ app = QApplication.instance()
812
+ if app and QThread.currentThread() is app.thread():
813
+ app.processEvents()
814
+
815
+ EPS = 1e-6
816
+
817
+ # -----------------------------
818
+ # Helpers: image prep / shapes
819
+ # -----------------------------
820
+
821
+ # new: lightweight loader that yields one frame at a time
822
+
823
+ def _to_luma_local(a: np.ndarray) -> np.ndarray:
824
+ a = np.asarray(a, dtype=np.float32)
825
+ if a.ndim == 2:
826
+ return a
827
+ if a.ndim == 3:
828
+ # mono fast paths
829
+ if a.shape[-1] == 1: # HWC mono
830
+ return a[..., 0].astype(np.float32, copy=False)
831
+ if a.shape[0] == 1: # CHW mono
832
+ return a[0].astype(np.float32, copy=False)
833
+ # RGB
834
+ if a.shape[-1] == 3: # HWC RGB
835
+ r, g, b = a[..., 0], a[..., 1], a[..., 2]
836
+ return (0.2126*r + 0.7152*g + 0.0722*b).astype(np.float32, copy=False)
837
+ if a.shape[0] == 3: # CHW RGB
838
+ r, g, b = a[0], a[1], a[2]
839
+ return (0.2126*r + 0.7152*g + 0.0722*b).astype(np.float32, copy=False)
840
+ # fallback: average last axis
841
+ return a.mean(axis=-1).astype(np.float32, copy=False)
842
+
843
+ def _normalize_layout_single(a, color_mode):
844
+ """
845
+ Coerce to:
846
+ - 'luma' -> (H, W)
847
+ - 'perchannel' -> (C, H, W); mono stays (1,H,W), RGB → (3,H,W)
848
+ Accepts (H,W), (H,W,3), or (3,H,W).
849
+ """
850
+ a = np.asarray(a, dtype=np.float32)
851
+
852
+ if color_mode == "luma":
853
+ return _to_luma_local(a) # returns (H,W)
854
+
855
+ # perchannel
856
+ if a.ndim == 2:
857
+ return a[None, ...] # (1,H,W) ← keep mono as 1 channel
858
+ if a.ndim == 3 and a.shape[-1] == 3:
859
+ return np.moveaxis(a, -1, 0) # (3,H,W)
860
+ if a.ndim == 3 and a.shape[0] in (1, 3):
861
+ return a # already (1,H,W) or (3,H,W)
862
+ # fallback: average any weird shape into luma 1×H×W
863
+ l = _to_luma_local(a)
864
+ return l[None, ...]
865
+
866
+
867
+ def _normalize_layout_batch(arrs, color_mode):
868
+ return [_normalize_layout_single(a, color_mode) for a in arrs]
869
+
870
+ def _common_hw(data_list):
871
+ """Return minimal (H,W) across items; items are (H,W) or (C,H,W)."""
872
+ Hs, Ws = [], []
873
+ for a in data_list:
874
+ if a.ndim == 2:
875
+ H, W = a.shape
876
+ else:
877
+ _, H, W = a.shape
878
+ Hs.append(H); Ws.append(W)
879
+ return int(min(Hs)), int(min(Ws))
880
+
881
+ def _center_crop(arr, Ht, Wt):
882
+ """Center-crop arr (H,W) or (C,H,W) to (Ht,Wt)."""
883
+ if arr.ndim == 2:
884
+ H, W = arr.shape
885
+ if H == Ht and W == Wt:
886
+ return arr
887
+ y0 = max(0, (H - Ht) // 2)
888
+ x0 = max(0, (W - Wt) // 2)
889
+ return arr[y0:y0+Ht, x0:x0+Wt]
890
+ else:
891
+ C, H, W = arr.shape
892
+ if H == Ht and W == Wt:
893
+ return arr
894
+ y0 = max(0, (H - Ht) // 2)
895
+ x0 = max(0, (W - Wt) // 2)
896
+ return arr[:, y0:y0+Ht, x0:x0+Wt]
897
+
898
+ def _sanitize_numeric(a):
899
+ """Replace NaN/Inf, clip negatives, make contiguous float32."""
900
+ a = np.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0)
901
+ a = np.clip(a, 0.0, None).astype(np.float32, copy=False)
902
+ return np.ascontiguousarray(a)
903
+
904
+ # -----------------------------
905
+ # PSF utilities
906
+ # -----------------------------
907
+
908
+ def _gaussian_psf(fwhm_px: float, ksize: int) -> np.ndarray:
909
+ sigma = max(fwhm_px, 1.0) / 2.3548
910
+ r = (ksize - 1) / 2
911
+ y, x = np.mgrid[-r:r+1, -r:r+1]
912
+ g = np.exp(-(x*x + y*y) / (2*sigma*sigma))
913
+ g /= (np.sum(g) + EPS)
914
+ return g.astype(np.float32, copy=False)
915
+
916
+ def _estimate_fwhm_from_header(hdr) -> float:
917
+ for key in ("FWHM", "FWHM_PIX", "PSF_FWHM"):
918
+ if key in hdr:
919
+ try:
920
+ val = float(hdr[key])
921
+ if np.isfinite(val) and val > 0:
922
+ return val
923
+ except Exception:
924
+ pass
925
+ return float("nan")
926
+
927
+ def _estimate_fwhm_from_image(arr) -> float:
928
+ """Fast FWHM estimate from SEP 'a','b' parameters (≈ sigma in px)."""
929
+ if sep is None:
930
+ return float("nan")
931
+ try:
932
+ img = _contig(_to_luma_local(arr)) # ← ensure C-contig float32
933
+ bkg = sep.Background(img)
934
+ data = _contig(img - bkg.back()) # ← ensure data is C-contig
935
+ try:
936
+ err = bkg.globalrms
937
+ except Exception:
938
+ err = float(np.median(bkg.rms()))
939
+ sources = sep.extract(data, 6.0, err=err)
940
+ if sources is None or len(sources) == 0:
941
+ return float("nan")
942
+ a = np.asarray(sources["a"], dtype=np.float32)
943
+ b = np.asarray(sources["b"], dtype=np.float32)
944
+ ab = (a + b) * 0.5
945
+ sigma = float(np.median(ab[np.isfinite(ab) & (ab > 0)]))
946
+ if not np.isfinite(sigma) or sigma <= 0:
947
+ return float("nan")
948
+ return 2.3548 * sigma
949
+ except Exception:
950
+ return float("nan")
951
+
952
+ def _auto_ksize_from_fwhm(fwhm_px: float, kmin: int = 11, kmax: int = 51) -> int:
953
+ """
954
+ Choose odd kernel size to cover about ±4σ.
955
+ """
956
+ sigma = max(fwhm_px, 1.0) / 2.3548
957
+ r = int(math.ceil(4.0 * sigma))
958
+ k = 2 * r + 1
959
+ k = max(kmin, min(k, kmax))
960
+ if (k % 2) == 0:
961
+ k += 1
962
+ return k
963
+
964
+ def _flip_kernel(psf):
965
+ # PyTorch dislikes negative strides; make it contiguous.
966
+ return np.flip(np.flip(psf, -1), -2).copy()
967
+
968
+ def _conv_same_np(img, psf):
969
+ # img: (H,W) or (C,H,W) numpy
970
+ import numpy.fft as fft
971
+ def fftconv2(a, k):
972
+ H, W = a.shape[-2:]
973
+ kh, kw = k.shape
974
+ pad_h, pad_w = H + kh - 1, W + kw - 1
975
+ A = fft.rfftn(a, s=(pad_h, pad_w), axes=(-2, -1))
976
+ K = fft.rfftn(k, s=(pad_h, pad_w), axes=(-2, -1))
977
+ Y = A * K
978
+ y = fft.irfftn(Y, s=(pad_h, pad_w), axes=(-2, -1))
979
+ sh, sw = (kh - 1)//2, (kw - 1)//2
980
+ return y[..., sh:sh+H, sw:sw+W]
981
+ if img.ndim == 2:
982
+ return fftconv2(img[None], psf)[0]
983
+ else:
984
+ return np.stack([fftconv2(img[c:c+1], psf)[0] for c in range(img.shape[0])], axis=0)
985
+
986
+ def _normalize_psf(psf):
987
+ psf = np.maximum(psf, 0.0).astype(np.float32, copy=False)
988
+ s = float(psf.sum())
989
+ if not np.isfinite(s) or s <= EPS:
990
+ return psf
991
+ return (psf / s).astype(np.float32, copy=False)
992
+
993
+ def _soften_psf(psf, sigma_px=0.25):
994
+ # optional tiny Gaussian soften to reduce ringing; sigma<=0 disables
995
+ if sigma_px <= 0:
996
+ return psf
997
+ r = int(max(1, round(3 * sigma_px)))
998
+ y, x = np.mgrid[-r:r+1, -r:r+1]
999
+ g = np.exp(-(x*x + y*y) / (2 * sigma_px * sigma_px)).astype(np.float32)
1000
+ g /= g.sum() + EPS
1001
+ return _conv_same_np(psf[None], g)[0]
1002
+
1003
+ def _psf_fwhm_px(psf: np.ndarray) -> float:
1004
+ """Approximate FWHM (pixels) from second moments of a normalized kernel."""
1005
+ psf = np.maximum(psf, 0).astype(np.float32, copy=False)
1006
+ s = float(psf.sum())
1007
+ if s <= EPS:
1008
+ return float("nan")
1009
+ k = psf.shape[0]
1010
+ y, x = np.mgrid[:k, :k].astype(np.float32)
1011
+ cy = float((psf * y).sum() / s)
1012
+ cx = float((psf * x).sum() / s)
1013
+ var_y = float((psf * (y - cy) ** 2).sum() / s)
1014
+ var_x = float((psf * (x - cx) ** 2).sum() / s)
1015
+ sigma = math.sqrt(max(0.0, 0.5 * (var_x + var_y)))
1016
+ return 2.3548 * sigma # FWHM≈2.355σ
1017
+
1018
+ STAR_MASK_MAXSIDE = 2048
1019
+ STAR_MASK_MAXOBJS = 2000 # cap number of objects
1020
+ VARMAP_SAMPLE_STRIDE = 8 # (kept for compat; currently unused internally)
1021
+ THRESHOLD_SIGMA = 2.0
1022
+ KEEP_FLOOR = 0.20
1023
+ GROW_PX = 8
1024
+ MAX_STAR_RADIUS = 16
1025
+ SOFT_SIGMA = 2.0
1026
+ ELLIPSE_SCALE = 1.2
1027
+
1028
+ def _sep_background_precompute(img_2d: np.ndarray, bw: int = 64, bh: int = 64):
1029
+ """One-time SEP background build; returns (sky_map, rms_map, err_scalar)."""
1030
+ if sep is None:
1031
+ # robust fallback
1032
+ med = float(np.median(img_2d))
1033
+ mad = float(np.median(np.abs(img_2d - med))) + 1e-6
1034
+ sky = np.full_like(img_2d, med, dtype=np.float32)
1035
+ rmsm = np.full_like(img_2d, 1.4826 * mad, dtype=np.float32)
1036
+ return sky, rmsm, float(np.median(rmsm))
1037
+
1038
+ a = np.ascontiguousarray(img_2d.astype(np.float32))
1039
+ b = sep.Background(a, bw=int(bw), bh=int(bh), fw=3, fh=3)
1040
+ sky = np.asarray(b.back(), dtype=np.float32)
1041
+ try:
1042
+ rmsm = np.asarray(b.rms(), dtype=np.float32)
1043
+ err = float(b.globalrms)
1044
+ except Exception:
1045
+ rmsm = np.full_like(a, float(np.median(b.rms())), dtype=np.float32)
1046
+ err = float(np.median(rmsm))
1047
+ return sky, rmsm, err
1048
+
1049
+
1050
+ def _star_mask_from_precomputed(
1051
+ img_2d: np.ndarray,
1052
+ sky_map: np.ndarray,
1053
+ err_scalar: float,
1054
+ *,
1055
+ thresh_sigma: float,
1056
+ max_objs: int,
1057
+ grow_px: int,
1058
+ ellipse_scale: float,
1059
+ soft_sigma: float,
1060
+ max_radius_px: int,
1061
+ keep_floor: float,
1062
+ max_side: int,
1063
+ status_cb=lambda s: None
1064
+ ) -> np.ndarray:
1065
+ """
1066
+ Build a KEEP weight map using a *downscaled detection / full-res draw* path.
1067
+ **Never writes to img_2d**; all drawing happens in a fresh `mask_u8`.
1068
+ """
1069
+ # Optional OpenCV fast path
1070
+ try:
1071
+ import cv2 as _cv2
1072
+ _HAS_CV2 = True
1073
+ except Exception:
1074
+ _HAS_CV2 = False
1075
+ _cv2 = None # type: ignore
1076
+
1077
+ H, W = map(int, img_2d.shape)
1078
+
1079
+ # Residual for detection (contiguous, separate buffer)
1080
+ data_sub = np.ascontiguousarray((img_2d - sky_map).astype(np.float32))
1081
+
1082
+ # Downscale *detection only* to speed up, never the draw step
1083
+ det = data_sub
1084
+ scale = 1.0
1085
+ if max_side and max(H, W) > int(max_side):
1086
+ scale = float(max(H, W)) / float(max_side)
1087
+ if _HAS_CV2:
1088
+ det = _cv2.resize(
1089
+ det,
1090
+ (max(1, int(round(W / scale))), max(1, int(round(H / scale)))),
1091
+ interpolation=_cv2.INTER_AREA
1092
+ )
1093
+ else:
1094
+ s = int(max(1, round(scale)))
1095
+ det = det[:(H // s) * s, :(W // s) * s].reshape(H // s, s, W // s, s).mean(axis=(1, 3))
1096
+ scale = float(s)
1097
+
1098
+ # Threshold ladder
1099
+ thresholds = [thresh_sigma, thresh_sigma*2, thresh_sigma*4,
1100
+ thresh_sigma*8, thresh_sigma*16]
1101
+ objs = None; used = float("nan"); raw = 0
1102
+ for t in thresholds:
1103
+ cand = sep.extract(det, thresh=float(t), err=float(err_scalar))
1104
+ n = 0 if cand is None else len(cand)
1105
+ if n == 0: continue
1106
+ if n > max_objs*12: continue
1107
+ objs, raw, used = cand, n, float(t)
1108
+ break
1109
+
1110
+ if objs is None or len(objs) == 0:
1111
+ try:
1112
+ cand = sep.extract(det, thresh=thresholds[-1], err=float(err_scalar), minarea=9)
1113
+ except Exception:
1114
+ cand = None
1115
+ if cand is None or len(cand) == 0:
1116
+ status_cb("Star mask: no sources found (mask disabled for this frame).")
1117
+ return np.ones((H, W), dtype=np.float32, order="C")
1118
+ objs, raw, used = cand, len(cand), float(thresholds[-1])
1119
+
1120
+ # Brightest max_objs
1121
+ if "flux" in objs.dtype.names:
1122
+ idx = np.argsort(objs["flux"])[-int(max_objs):]
1123
+ objs = objs[idx]
1124
+ else:
1125
+ objs = objs[:int(max_objs)]
1126
+ kept = len(objs)
1127
+
1128
+ # ---- draw back on full-res into a brand-new buffer ----
1129
+ mask_u8 = np.zeros((H, W), dtype=np.uint8, order="C")
1130
+ s_back = float(scale)
1131
+ MR = int(max(1, max_radius_px))
1132
+ G = int(max(0, grow_px))
1133
+ ES = float(max(0.1, ellipse_scale))
1134
+
1135
+ drawn = 0
1136
+ if _HAS_CV2:
1137
+ for o in objs:
1138
+ x = int(round(float(o["x"]) * s_back))
1139
+ y = int(round(float(o["y"]) * s_back))
1140
+ if not (0 <= x < W and 0 <= y < H):
1141
+ continue
1142
+ a = float(o["a"]) * s_back
1143
+ b = float(o["b"]) * s_back
1144
+ r = int(math.ceil(ES * max(a, b)))
1145
+ r = min(max(r, 0) + G, MR)
1146
+ if r <= 0:
1147
+ continue
1148
+ _cv2.circle(mask_u8, (x, y), r, 1, thickness=-1, lineType=_cv2.LINE_8)
1149
+ drawn += 1
1150
+ else:
1151
+ for o in objs:
1152
+ x = int(round(float(o["x"]) * s_back))
1153
+ y = int(round(float(o["y"]) * s_back))
1154
+ if not (0 <= x < W and 0 <= y < H):
1155
+ continue
1156
+ a = float(o["a"]) * s_back
1157
+ b = float(o["b"]) * s_back
1158
+ r = int(math.ceil(ES * max(a, b)))
1159
+ r = min(max(r, 0) + G, MR)
1160
+ if r <= 0:
1161
+ continue
1162
+ y0 = max(0, y - r); y1 = min(H, y + r + 1)
1163
+ x0 = max(0, x - r); x1 = min(W, x + r + 1)
1164
+ yy, xx = np.ogrid[y0:y1, x0:x1]
1165
+ disk = (yy - y)*(yy - y) + (xx - x)*(xx - x) <= r*r
1166
+ mask_u8[y0:y1, x0:x1][disk] = 1
1167
+ drawn += 1
1168
+
1169
+ # Feather + convert to keep weights
1170
+ m = mask_u8.astype(np.float32, copy=False)
1171
+ if soft_sigma > 0:
1172
+ try:
1173
+ if _HAS_CV2:
1174
+ k = int(max(1, int(round(3*soft_sigma)))*2 + 1)
1175
+ m = _cv2.GaussianBlur(m, (k, k), float(soft_sigma),
1176
+ borderType=_cv2.BORDER_REFLECT)
1177
+ else:
1178
+ from scipy.ndimage import gaussian_filter
1179
+ m = gaussian_filter(m, sigma=float(soft_sigma), mode="reflect")
1180
+ except Exception:
1181
+ pass
1182
+ np.clip(m, 0.0, 1.0, out=m)
1183
+
1184
+ keep = 1.0 - m
1185
+ kf = float(max(0.0, min(0.99, keep_floor)))
1186
+ keep = kf + (1.0 - kf) * keep
1187
+ np.clip(keep, 0.0, 1.0, out=keep)
1188
+
1189
+ status_cb(f"Star mask: thresh={used:.3g} | detected={raw} | kept={kept} | drawn={drawn} | keep_floor={keep_floor}")
1190
+ return np.ascontiguousarray(keep, dtype=np.float32)
1191
+
1192
+
1193
+ def _variance_map_from_precomputed(
1194
+ img_2d: np.ndarray,
1195
+ sky_map: np.ndarray,
1196
+ rms_map: np.ndarray,
1197
+ hdr,
1198
+ *,
1199
+ smooth_sigma: float,
1200
+ floor: float,
1201
+ status_cb=lambda s: None
1202
+ ) -> np.ndarray:
1203
+ img = np.clip(np.asarray(img_2d, dtype=np.float32), 0.0, None)
1204
+ var_bg_dn2 = np.maximum(rms_map, 1e-6) ** 2
1205
+ obj_dn = np.clip(img - sky_map, 0.0, None)
1206
+
1207
+ gain = None
1208
+ for k in ("EGAIN", "GAIN", "GAIN1", "GAIN2"):
1209
+ if k in hdr:
1210
+ try:
1211
+ g = float(hdr[k]); gain = g if (np.isfinite(g) and g > 0) else None
1212
+ if gain is not None: break
1213
+ except Exception as e:
1214
+ import logging
1215
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
1216
+
1217
+ if gain is not None:
1218
+ a_shot = 1.0 / gain
1219
+ else:
1220
+ sky_med = float(np.median(sky_map))
1221
+ varbg_med= float(np.median(var_bg_dn2))
1222
+ a_shot = (varbg_med / sky_med) if sky_med > 1e-6 else 0.0
1223
+ a_shot = float(np.clip(a_shot, 0.0, 10.0))
1224
+
1225
+ v = var_bg_dn2 + a_shot * obj_dn
1226
+ if smooth_sigma > 0:
1227
+ try:
1228
+ import cv2 as _cv2
1229
+ k = int(max(1, int(round(3*smooth_sigma)))*2 + 1)
1230
+ v = _cv2.GaussianBlur(v, (k,k), float(smooth_sigma), borderType=_cv2.BORDER_REFLECT)
1231
+ except Exception:
1232
+ try:
1233
+ from scipy.ndimage import gaussian_filter
1234
+ v = gaussian_filter(v, sigma=float(smooth_sigma), mode="reflect")
1235
+ except Exception:
1236
+ pass
1237
+
1238
+ np.clip(v, float(floor), None, out=v)
1239
+ try:
1240
+ rms_med = float(np.median(np.sqrt(var_bg_dn2)))
1241
+ status_cb(f"Variance map: sky_med={float(np.median(sky_map)):.3g} DN | rms_med={rms_med:.3g} DN | smooth_sigma={smooth_sigma} | floor={floor}")
1242
+ except Exception:
1243
+ pass
1244
+ return v.astype(np.float32, copy=False)
1245
+
1246
+
1247
+
1248
+ # -----------------------------
1249
+ # Robust weighting (Huber)
1250
+ # -----------------------------
1251
+
1252
+ def _estimate_scalar_variance_t(r):
1253
+ # r: tensor on device
1254
+ med = torch.median(r)
1255
+ mad = torch.median(torch.abs(r - med)) + 1e-6
1256
+ return (1.4826 * mad) ** 2
1257
+
1258
+ def _estimate_scalar_variance(a):
1259
+ med = np.median(a)
1260
+ mad = np.median(np.abs(a - med)) + 1e-6
1261
+ return float((1.4826 * mad) ** 2)
1262
+
1263
+ def _weight_map(y, pred, huber_delta, var_map=None, mask=None):
1264
+ """
1265
+ Robust per-pixel weights for the MM update.
1266
+ W = [psi(r)/r] * 1/(var + eps) * mask
1267
+ If huber_delta < 0, delta = (-huber_delta) * RMS(residual) (auto).
1268
+ var_map: per-pixel variance (2D); if None, fall back to robust scalar via MAD.
1269
+ mask: 2D {0,1} validity; if None, treat as ones.
1270
+ """
1271
+ r = y - pred
1272
+ eps = EPS
1273
+
1274
+ # resolve Huber delta
1275
+ if huber_delta < 0:
1276
+ if TORCH_OK and isinstance(r, torch.Tensor):
1277
+ med = torch.median(r)
1278
+ mad = torch.median(torch.abs(r - med)) + 1e-6
1279
+ rms = 1.4826 * mad
1280
+ delta = (-huber_delta) * torch.clamp(rms, min=1e-6)
1281
+ else:
1282
+ med = np.median(r)
1283
+ mad = np.median(np.abs(r - med)) + 1e-6
1284
+ rms = 1.4826 * mad
1285
+ delta = (-huber_delta) * max(rms, 1e-6)
1286
+ else:
1287
+ delta = huber_delta
1288
+
1289
+ # psi(r)/r
1290
+ if TORCH_OK and isinstance(r, torch.Tensor):
1291
+ absr = torch.abs(r)
1292
+ if float(delta) > 0:
1293
+ psi_over_r = torch.where(absr <= delta, torch.ones_like(r), delta / (absr + eps))
1294
+ else:
1295
+ psi_over_r = torch.ones_like(r)
1296
+ if var_map is None:
1297
+ v = _estimate_scalar_variance_t(r)
1298
+ else:
1299
+ v = var_map
1300
+ if v.ndim == 2 and r.ndim == 3:
1301
+ v = v[None, ...] # broadcast over channels
1302
+ w = psi_over_r / (v + eps)
1303
+ if mask is not None:
1304
+ m = mask if mask.ndim == w.ndim else (mask[None, ...] if w.ndim == 3 else mask)
1305
+ w = w * m
1306
+ return w
1307
+ else:
1308
+ absr = np.abs(r)
1309
+ if float(delta) > 0:
1310
+ psi_over_r = np.where(absr <= delta, 1.0, delta / (absr + eps)).astype(np.float32)
1311
+ else:
1312
+ psi_over_r = np.ones_like(r, dtype=np.float32)
1313
+ if var_map is None:
1314
+ v = _estimate_scalar_variance(r)
1315
+ else:
1316
+ v = var_map
1317
+ if v.ndim == 2 and r.ndim == 3:
1318
+ v = v[None, ...]
1319
+ w = psi_over_r / (v + eps)
1320
+ if mask is not None:
1321
+ m = mask if mask.ndim == w.ndim else (mask[None, ...] if w.ndim == 3 else mask)
1322
+ w = w * m
1323
+ return w
1324
+
1325
+
1326
+ # -----------------------------
1327
+ # Torch / conv
1328
+ # -----------------------------
1329
+
1330
+ def _fftshape_same(H, W, kh, kw):
1331
+ return H + kh - 1, W + kw - 1
1332
+
1333
+ # ---------- Torch FFT helpers (FIXED: carry padH/padW) ----------
1334
+ def _precompute_torch_psf_ffts(psfs, flip_psf, H, W, device, dtype):
1335
+ tfft = torch.fft
1336
+ psf_fft, psfT_fft = [], []
1337
+ for k, kT in zip(psfs, flip_psf):
1338
+ kh, kw = k.shape
1339
+ padH, padW = _fftshape_same(H, W, kh, kw)
1340
+
1341
+ # shift the small kernels to the origin, then FFT into padded size
1342
+ k_small = torch.as_tensor(np.fft.ifftshift(k), device=device, dtype=dtype)
1343
+ kT_small = torch.as_tensor(np.fft.ifftshift(kT), device=device, dtype=dtype)
1344
+
1345
+ Kf = tfft.rfftn(k_small, s=(padH, padW))
1346
+ KTf = tfft.rfftn(kT_small, s=(padH, padW))
1347
+
1348
+ psf_fft.append((Kf, padH, padW, kh, kw))
1349
+ psfT_fft.append((KTf, padH, padW, kh, kw))
1350
+ return psf_fft, psfT_fft
1351
+
1352
+
1353
+
1354
+ # ---------- NumPy FFT helpers ----------
1355
+ def _precompute_np_psf_ffts(psfs, flip_psf, H, W):
1356
+ import numpy.fft as fft
1357
+ meta, Kfs, KTfs = [], [], []
1358
+ for k, kT in zip(psfs, flip_psf):
1359
+ kh, kw = k.shape
1360
+ fftH, fftW = _fftshape_same(H, W, kh, kw)
1361
+ Kfs.append( fft.rfftn(np.fft.ifftshift(k), s=(fftH, fftW)) )
1362
+ KTfs.append(fft.rfftn(np.fft.ifftshift(kT), s=(fftH, fftW)) )
1363
+ meta.append((kh, kw, fftH, fftW))
1364
+ return Kfs, KTfs, meta
1365
+
1366
+ def _fft_conv_same_np(a, Kf, kh, kw, fftH, fftW, out):
1367
+ import numpy.fft as fft
1368
+ if a.ndim == 2:
1369
+ A = fft.rfftn(a, s=(fftH, fftW))
1370
+ y = fft.irfftn(A * Kf, s=(fftH, fftW))
1371
+ sh, sw = kh // 2, kw // 2
1372
+ out[...] = y[sh:sh+a.shape[0], sw:sw+a.shape[1]]
1373
+ return out
1374
+ else:
1375
+ C, H, W = a.shape
1376
+ acc = []
1377
+ for c in range(C):
1378
+ A = fft.rfftn(a[c], s=(fftH, fftW))
1379
+ y = fft.irfftn(A * Kf, s=(fftH, fftW))
1380
+ sh, sw = kh // 2, kw // 2
1381
+ acc.append(y[sh:sh+H, sw:sw+W])
1382
+ out[...] = np.stack(acc, 0)
1383
+ return out
1384
+
1385
+
1386
+
1387
+ def _torch_device():
1388
+ if TORCH_OK and (torch is not None):
1389
+ if hasattr(torch, "cuda") and torch.cuda.is_available():
1390
+ return torch.device("cuda")
1391
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
1392
+ return torch.device("mps")
1393
+ # DirectML: we passed dml_device from outer scope; keep a module-global
1394
+ if globals().get("dml_ok", False) and globals().get("dml_device", None) is not None:
1395
+ return globals()["dml_device"]
1396
+ return torch.device("cpu")
1397
+
1398
+ def _to_t(x: np.ndarray):
1399
+ if not (TORCH_OK and (torch is not None)):
1400
+ raise RuntimeError("Torch path requested but torch is unavailable")
1401
+ device = _torch_device()
1402
+ t = torch.from_numpy(x)
1403
+ # DirectML wants explicit .to(device)
1404
+ return t.to(device, non_blocking=True) if str(device) != "cpu" else t
1405
+
1406
+ def _contig(x):
1407
+ return np.ascontiguousarray(x, dtype=np.float32)
1408
+
1409
+ def _conv_same_torch(img_t, psf_t):
1410
+ """
1411
+ img_t: torch tensor on DEVICE, (H,W) or (C,H,W)
1412
+ psf_t: torch tensor on DEVICE, (1,1,kh,kw) (single kernel)
1413
+ Pads with 'reflect' to avoid zero-padding ringing.
1414
+ """
1415
+ kh, kw = psf_t.shape[-2:]
1416
+ pad = (kw // 2, kw - kw // 2 - 1, # left, right
1417
+ kh // 2, kh - kh // 2 - 1) # top, bottom
1418
+
1419
+ if img_t.ndim == 2:
1420
+ x = img_t[None, None]
1421
+ x = torch.nn.functional.pad(x, pad, mode="reflect")
1422
+ y = torch.nn.functional.conv2d(x, psf_t, padding=0)
1423
+ return y[0, 0]
1424
+ else:
1425
+ C = img_t.shape[0]
1426
+ x = img_t[None]
1427
+ x = torch.nn.functional.pad(x, pad, mode="reflect")
1428
+ w = psf_t.repeat(C, 1, 1, 1)
1429
+ y = torch.nn.functional.conv2d(x, w, padding=0, groups=C)
1430
+ return y[0]
1431
+
1432
+ def _safe_inference_context():
1433
+ """
1434
+ Return a valid, working no-grad context:
1435
+ - prefer torch.inference_mode() if it exists *and* can be entered,
1436
+ - otherwise fall back to torch.no_grad(),
1437
+ - if torch is unavailable, return NO_GRAD.
1438
+ """
1439
+ if not (TORCH_OK and (torch is not None)):
1440
+ return NO_GRAD
1441
+
1442
+ cm = getattr(torch, "inference_mode", None)
1443
+ if cm is None:
1444
+ return torch.no_grad
1445
+
1446
+ # Probe inference_mode once; if it explodes on this build, fall back.
1447
+ try:
1448
+ with cm():
1449
+ pass
1450
+ return cm
1451
+ except Exception:
1452
+ return torch.no_grad
1453
+
1454
+ def _ensure_mask_list(masks, data):
1455
+ # 1s where valid, 0s where invalid (soft edges allowed)
1456
+ if masks is None:
1457
+ return [np.ones_like(a if a.ndim==2 else a[0], dtype=np.float32) for a in data]
1458
+ out = []
1459
+ for a, m in zip(data, masks):
1460
+ base = a if a.ndim==2 else a[0] # mask is 2D; shared across channels
1461
+ if m is None:
1462
+ out.append(np.ones_like(base, dtype=np.float32))
1463
+ else:
1464
+ mm = np.asarray(m, dtype=np.float32)
1465
+ if mm.ndim == 3: # tolerate (1,H,W) or (C,H,W)
1466
+ mm = mm[0]
1467
+ if mm.shape != base.shape:
1468
+ # center crop to match (common intersection already applied)
1469
+ Ht, Wt = base.shape
1470
+ mm = _center_crop(mm, Ht, Wt)
1471
+ # keep as float weights in [0,1] (do not threshold!)
1472
+ out.append(np.clip(mm.astype(np.float32, copy=False), 0.0, 1.0))
1473
+ return out
1474
+
1475
+ def _ensure_var_list(variances, data):
1476
+ # If None, we’ll estimate a robust scalar per frame on-the-fly.
1477
+ if variances is None:
1478
+ return [None]*len(data)
1479
+ out = []
1480
+ for a, v in zip(data, variances):
1481
+ if v is None:
1482
+ out.append(None)
1483
+ else:
1484
+ vv = np.asarray(v, dtype=np.float32)
1485
+ if vv.ndim == 3:
1486
+ vv = vv[0]
1487
+ base = a if a.ndim==2 else a[0]
1488
+ if vv.shape != base.shape:
1489
+ Ht, Wt = base.shape
1490
+ vv = _center_crop(vv, Ht, Wt)
1491
+ # clip tiny/negatives
1492
+ vv = np.clip(vv, 1e-8, None).astype(np.float32, copy=False)
1493
+ out.append(vv)
1494
+ return out
1495
+
1496
+ # ---- SR operators (downsample / upsample-sum) ----
1497
+ def _downsample_avg(img, r: int):
1498
+ """Average-pool over non-overlapping r×r blocks. Works for (H,W) or (C,H,W)."""
1499
+ if r <= 1:
1500
+ return img
1501
+ a = np.asarray(img, dtype=np.float32)
1502
+ if a.ndim == 2:
1503
+ H, W = a.shape
1504
+ Hs, Ws = (H // r) * r, (W // r) * r
1505
+ a = a[:Hs, :Ws].reshape(Hs//r, r, Ws//r, r).mean(axis=(1,3))
1506
+ return a
1507
+ else:
1508
+ C, H, W = a.shape
1509
+ Hs, Ws = (H // r) * r, (W // r) * r
1510
+ a = a[:, :Hs, :Ws].reshape(C, Hs//r, r, Ws//r, r).mean(axis=(2,4))
1511
+ return a
1512
+
1513
+ def _upsample_sum(img, r: int, target_hw: tuple[int,int] | None = None):
1514
+ """Adjoint of average-pooling: replicate-sum each pixel into an r×r block.
1515
+ For (H,W) or (C,H,W). If target_hw given, center-crop/pad to that size.
1516
+ """
1517
+ if r <= 1:
1518
+ return img
1519
+ a = np.asarray(img, dtype=np.float32)
1520
+ if a.ndim == 2:
1521
+ H, W = a.shape
1522
+ out = np.kron(a, np.ones((r, r), dtype=np.float32))
1523
+ else:
1524
+ C, H, W = a.shape
1525
+ out = np.stack([np.kron(a[c], np.ones((r, r), dtype=np.float32)) for c in range(C)], axis=0)
1526
+ if target_hw is not None:
1527
+ Ht, Wt = target_hw
1528
+ out = _center_crop(out, Ht, Wt)
1529
+ return out
1530
+
1531
+ def _gaussian2d(ksize: int, sigma: float) -> np.ndarray:
1532
+ r = (ksize - 1) // 2
1533
+ y, x = np.mgrid[-r:r+1, -r:r+1].astype(np.float32)
1534
+ g = np.exp(-(x*x + y*y)/(2.0*sigma*sigma)).astype(np.float32)
1535
+ g /= g.sum() + EPS
1536
+ return g
1537
+
1538
+ def _conv2_same_np(a: np.ndarray, k: np.ndarray) -> np.ndarray:
1539
+ # lightweight wrap for 2D conv on (H,W) or (C,H,W) with same-size output
1540
+ return _conv_same_np(a if a.ndim==3 else a[None], k)[0] if a.ndim==2 else _conv_same_np(a, k)
1541
+
1542
+ def _solve_super_psf_from_native(f_native: np.ndarray, r: int, sigma: float = 1.1,
1543
+ iters: int = 500, lr: float = 0.1) -> np.ndarray:
1544
+ """
1545
+ Solve: h* = argmin_h || f_native - (D(h) * g_sigma) ||_2^2,
1546
+ where h is (r*k)×(r*k) if f_native is k×k. Returns normalized h (sum=1).
1547
+ """
1548
+ f = np.asarray(f_native, dtype=np.float32)
1549
+ k = int(f.shape[0]); assert f.shape[0] == f.shape[1]
1550
+ kr = int(k * r)
1551
+
1552
+ # build Gaussian pre-blur at native scale (match paper §4.2)
1553
+ g = _gaussian2d(k, max(sigma, 1e-3)).astype(np.float32)
1554
+
1555
+ # init h by zero-insertion (nearest upsample of f) then deconvolving g very mildly
1556
+ h0 = np.zeros((kr, kr), dtype=np.float32)
1557
+ h0[::r, ::r] = f
1558
+ h0 = _normalize_psf(h0)
1559
+
1560
+ if TORCH_OK:
1561
+ dev = _torch_device()
1562
+ t = torch.tensor(h0, device=dev, dtype=torch.float32, requires_grad=True)
1563
+ f_t = torch.tensor(f, device=dev, dtype=torch.float32)
1564
+ g_t = torch.tensor(g, device=dev, dtype=torch.float32)
1565
+ opt = torch.optim.Adam([t], lr=lr)
1566
+ for _ in range(max(10, iters)):
1567
+ opt.zero_grad(set_to_none=True)
1568
+ H, W = t.shape
1569
+ Hr, Wr = H//r, W//r
1570
+ th = t[:Hr*r, :Wr*r].reshape(Hr, r, Wr, r).mean(dim=(1,3))
1571
+ # conv native: (Dh) * g
1572
+ conv = torch.nn.functional.conv2d(th[None,None], g_t[None,None], padding=g_t.shape[-1]//2)[0,0]
1573
+ loss = torch.mean((conv - f_t)**2)
1574
+ loss.backward()
1575
+ opt.step()
1576
+ with torch.no_grad():
1577
+ t.clamp_(min=0.0)
1578
+ t /= (t.sum() + 1e-8)
1579
+ h = t.detach().cpu().numpy().astype(np.float32)
1580
+ else:
1581
+ # Tiny gradient-descent fallback on numpy
1582
+ h = h0.copy()
1583
+ eta = float(lr)
1584
+ for _ in range(max(50, iters)):
1585
+ Dh = _downsample_avg(h, r)
1586
+ conv = _conv2_same_np(Dh, g)
1587
+ resid = (conv - f)
1588
+ # backprop through conv and D: grad wrt Dh is resid * g^T conv; adjoint of D is upsample-sum
1589
+ grad_Dh = _conv2_same_np(resid, np.flip(np.flip(g, 0), 1))
1590
+ grad_h = _upsample_sum(grad_Dh, r, target_hw=h.shape)
1591
+ h = np.clip(h - eta * grad_h, 0.0, None)
1592
+ s = float(h.sum()); h /= (s + 1e-8)
1593
+ eta *= 0.995
1594
+ return _normalize_psf(h)
1595
+
1596
+ def _downsample_avg_t(x, r: int):
1597
+ """
1598
+ Average-pool over non-overlapping r×r blocks.
1599
+ Works for (H,W) or (C,H,W). Crops to multiples of r.
1600
+ """
1601
+ if r <= 1:
1602
+ return x
1603
+ if x.ndim == 2:
1604
+ H, W = x.shape
1605
+ Hr, Wr = (H // r) * r, (W // r) * r
1606
+ if Hr == 0 or Wr == 0:
1607
+ return x # nothing to pool
1608
+ x2 = x[:Hr, :Wr]
1609
+ return x2.view(Hr // r, r, Wr // r, r).mean(dim=(1, 3))
1610
+ else:
1611
+ C, H, W = x.shape
1612
+ Hr, Wr = (H // r) * r, (W // r) * r
1613
+ if Hr == 0 or Wr == 0:
1614
+ return x
1615
+ x2 = x[:, :Hr, :Wr]
1616
+ return x2.view(C, Hr // r, r, Wr // r, r).mean(dim=(2, 4))
1617
+
1618
+ def _upsample_sum_t(x, r: int):
1619
+ if r <= 1:
1620
+ return x
1621
+ if x.ndim == 2:
1622
+ return x.repeat_interleave(r, dim=0).repeat_interleave(r, dim=1)
1623
+ else:
1624
+ return x.repeat_interleave(r, dim=-2).repeat_interleave(r, dim=-1)
1625
+
1626
+ def _sep_bg_rms(frames):
1627
+ """Return a robust background RMS using SEP's background model on the first frame."""
1628
+ if sep is None or not frames:
1629
+ return None
1630
+ try:
1631
+ y0 = frames[0] if frames[0].ndim == 2 else frames[0][0] # use luma/first channel
1632
+ a = np.ascontiguousarray(y0, dtype=np.float32)
1633
+ b = sep.Background(a, bw=64, bh=64, fw=3, fh=3)
1634
+ try:
1635
+ rms_val = float(b.globalrms)
1636
+ except Exception:
1637
+ # some SEP builds don’t expose globalrms; fall back to the map’s median
1638
+ rms_val = float(np.median(np.asarray(b.rms(), dtype=np.float32)))
1639
+ return rms_val
1640
+ except Exception:
1641
+ return None
1642
+
1643
+ # =========================
1644
+ # Memory/streaming helpers
1645
+ # =========================
1646
+
1647
+ def _approx_bytes(arr_like_shape, dtype=np.float32):
1648
+ """Rough byte estimator for a given shape/dtype."""
1649
+ return int(np.prod(arr_like_shape)) * np.dtype(dtype).itemsize
1650
+
1651
+
1652
+
1653
+ def _read_shape_fast(path) -> tuple[int,int,int]:
1654
+ if _is_xisf(path):
1655
+ a, _ = _load_image_array(path)
1656
+ if a is None:
1657
+ raise ValueError(f"No data in {path}")
1658
+ a = np.asarray(a)
1659
+ else:
1660
+ with fits.open(path, memmap=True, ignore_missing_simple=True) as hdul:
1661
+ a = hdul[0].data
1662
+ if a is None:
1663
+ raise ValueError(f"No data in {path}")
1664
+
1665
+ # common logic for both XISF and FITS
1666
+ if a.ndim == 2:
1667
+ H, W = a.shape
1668
+ return (1, int(H), int(W))
1669
+ if a.ndim == 3:
1670
+ if a.shape[-1] in (1, 3): # HWC
1671
+ C = int(a.shape[-1]); H = int(a.shape[0]); W = int(a.shape[1])
1672
+ return (1 if C == 1 else 3, H, W)
1673
+ if a.shape[0] in (1, 3): # CHW
1674
+ return (int(a.shape[0]), int(a.shape[1]), int(a.shape[2]))
1675
+ s = tuple(map(int, a.shape))
1676
+ H, W = s[-2], s[-1]
1677
+ return (1, H, W)
1678
+
1679
+
1680
+ def _read_tile_fits_any(path: str, y0: int, y1: int, x0: int, x1: int) -> np.ndarray:
1681
+ """FITS/XISF-aware tile read: returns spatial tile; supports 2D, HWC, and CHW."""
1682
+ ext = os.path.splitext(path)[1].lower()
1683
+
1684
+ if ext == ".xisf":
1685
+ a, _ = _load_image_array(path) # helper returns array-like + hdr/metadata
1686
+ if a is None:
1687
+ raise ValueError(f"XISF loader returned None for {path}")
1688
+ a = np.asarray(a)
1689
+ if a.ndim == 2: # HW
1690
+ return np.array(a[y0:y1, x0:x1], copy=True)
1691
+ elif a.ndim == 3:
1692
+ if a.shape[-1] in (1, 3): # HWC
1693
+ out = a[y0:y1, x0:x1, :]
1694
+ if out.shape[-1] == 1:
1695
+ out = out[..., 0]
1696
+ return np.array(out, copy=True)
1697
+ elif a.shape[0] in (1, 3): # CHW
1698
+ out = a[:, y0:y1, x0:x1]
1699
+ if out.shape[0] == 1:
1700
+ out = out[0]
1701
+ return np.array(out, copy=True)
1702
+ else:
1703
+ raise ValueError(f"Unsupported XISF 3D shape {a.shape} in {path}")
1704
+ else:
1705
+ raise ValueError(f"Unsupported XISF ndim {a.ndim} in {path}")
1706
+
1707
+ # FITS
1708
+ with fits.open(path, memmap=True, ignore_missing_simple=True) as hdul:
1709
+ a = None
1710
+ for h in hdul:
1711
+ if getattr(h, "data", None) is not None:
1712
+ a = h.data
1713
+ break
1714
+ if a is None:
1715
+ raise ValueError(f"No image data in {path}")
1716
+
1717
+ a = np.asarray(a)
1718
+
1719
+ if a.ndim == 2: # HW
1720
+ return np.array(a[y0:y1, x0:x1], copy=True)
1721
+
1722
+ if a.ndim == 3:
1723
+ if a.shape[0] in (1, 3): # CHW (planes, rows, cols)
1724
+ out = a[:, y0:y1, x0:x1]
1725
+ if out.shape[0] == 1: out = out[0]
1726
+ return np.array(out, copy=True)
1727
+ if a.shape[-1] in (1, 3): # HWC
1728
+ out = a[y0:y1, x0:x1, :]
1729
+ if out.shape[-1] == 1: out = out[..., 0]
1730
+ return np.array(out, copy=True)
1731
+
1732
+ # Fallback: assume last two axes are spatial (…, H, W)
1733
+ try:
1734
+ out = a[(..., slice(y0, y1), slice(x0, x1))]
1735
+ return np.array(out, copy=True)
1736
+ except Exception:
1737
+ raise ValueError(f"Unsupported FITS data shape {a.shape} in {path}")
1738
+
1739
+
1740
+ def _seed_median_full_from_data(data_list):
1741
+ """
1742
+ data_list: list of np.ndarray each shaped either (H,W) or (C,H,W),
1743
+ already cropped/sanitized to the same size by the caller.
1744
+ Returns: (H,W) or (C,H,W) median image in float32.
1745
+ """
1746
+ if not data_list:
1747
+ raise ValueError("Empty stack for median seed")
1748
+
1749
+ a0 = data_list[0]
1750
+ if a0.ndim == 2:
1751
+ # (N, H, W) -> (H, W)
1752
+ cube = np.stack([np.asarray(a, dtype=np.float32, order="C") for a in data_list], axis=0)
1753
+ med = np.median(cube, axis=0).astype(np.float32, copy=False)
1754
+ return np.ascontiguousarray(med)
1755
+ else:
1756
+ # (N, C, H, W) -> (C, H, W)
1757
+ cube = np.stack([np.asarray(a, dtype=np.float32, order="C") for a in data_list], axis=0)
1758
+ med = np.median(cube, axis=0).astype(np.float32, copy=False)
1759
+ return np.ascontiguousarray(med)
1760
+
1761
+
1762
+ def _build_seed_running_mu_sigma_from_paths(paths, Ht, Wt, color_mode,
1763
+ *, bootstrap_frames=20, clip_sigma=5.0,
1764
+ status_cb=lambda s: None, progress_cb=lambda f,m='': None):
1765
+ K = max(1, min(int(bootstrap_frames), len(paths)))
1766
+ def _load_chw(i):
1767
+ ys, _ = _stack_loader_memmap([paths[i]], Ht, Wt, color_mode)
1768
+ return _as_chw(ys[0]).astype(np.float32, copy=False)
1769
+ x0 = _load_chw(0).copy()
1770
+ mean = x0; m2 = np.zeros_like(mean); count = 1
1771
+ for i in range(1, K):
1772
+ x = _load_chw(i); count += 1
1773
+ d = x - mean; mean += d/count; m2 += d*(x-mean)
1774
+ progress_cb(i/K*0.5, "μ-σ bootstrap")
1775
+ var = m2 / max(1, count-1); sigma = np.sqrt(np.clip(var, 1e-12, None)).astype(np.float32)
1776
+ lo = mean - float(clip_sigma)*sigma; hi = mean + float(clip_sigma)*sigma
1777
+ acc = np.zeros_like(mean); n=0
1778
+ for i in range(len(paths)):
1779
+ x = _load_chw(i); x = np.clip(x, lo, hi, out=x)
1780
+ acc += x; n += 1; progress_cb(0.5 + 0.5*(i+1)/len(paths), "clipped mean")
1781
+ seed = (acc/max(1,n)).astype(np.float32)
1782
+ return seed[0] if (seed.ndim==3 and seed.shape[0]==1) else seed
1783
+
1784
+ # -----------------------------
1785
+ # Core
1786
+ # -----------------------------
1787
+ def multiframe_deconv(
1788
+ paths,
1789
+ out_path,
1790
+ iters=20,
1791
+ kappa=2.0,
1792
+ color_mode="luma",
1793
+ seed_mode: str = "robust",
1794
+ huber_delta=0.0,
1795
+ masks=None,
1796
+ variances=None,
1797
+ rho="huber",
1798
+ status_cb=lambda s: None,
1799
+ min_iters: int = 3,
1800
+ use_star_masks: bool = False,
1801
+ use_variance_maps: bool = False,
1802
+ star_mask_cfg: dict | None = None,
1803
+ varmap_cfg: dict | None = None,
1804
+ save_intermediate: bool = False,
1805
+ save_every: int = 1,
1806
+ # >>> SR options
1807
+ super_res_factor: int = 1,
1808
+ sr_sigma: float = 1.1,
1809
+ sr_psf_opt_iters: int = 250,
1810
+ sr_psf_opt_lr: float = 0.1,
1811
+ star_mask_ref_path: str | None = None,
1812
+ ):
1813
+ # sanitize and clamp
1814
+ max_iters = max(1, int(iters))
1815
+ min_iters = max(1, int(min_iters))
1816
+ if min_iters > max_iters:
1817
+ min_iters = max_iters
1818
+
1819
+ def _emit_pct(pct: float, msg: str | None = None):
1820
+ pct = float(max(0.0, min(1.0, pct)))
1821
+ status_cb(f"__PROGRESS__ {pct:.4f}" + (f" {msg}" if msg else ""))
1822
+
1823
+ status_cb(f"MFDeconv: loading {len(paths)} aligned frames…")
1824
+ _emit_pct(0.02, "loading")
1825
+
1826
+ # Use unified probe to pick a common crop without loading full images
1827
+ Ht, Wt = _common_hw_from_paths(paths)
1828
+ _emit_pct(0.05, "preparing")
1829
+
1830
+ # Stream actual pixels cropped to (Ht,Wt), float32 CHW/2D + headers
1831
+ ys_raw, hdrs = _stack_loader_memmap(paths, Ht, Wt, color_mode)
1832
+ relax = 0.7
1833
+ use_torch = False
1834
+ global torch, TORCH_OK
1835
+
1836
+ # -------- try to import torch from per-user runtime venv --------
1837
+ torch = None
1838
+ TORCH_OK = False
1839
+ cuda_ok = mps_ok = dml_ok = False
1840
+ dml_device = None
1841
+ try:
1842
+ from setiastro.saspro.runtime_torch import import_torch
1843
+ torch = import_torch(prefer_cuda=True, status_cb=status_cb)
1844
+ TORCH_OK = True
1845
+
1846
+ try: cuda_ok = hasattr(torch, "cuda") and torch.cuda.is_available()
1847
+ except Exception: cuda_ok = False
1848
+ try: mps_ok = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
1849
+ except Exception: mps_ok = False
1850
+ try:
1851
+ import torch_directml
1852
+ dml_device = torch_directml.device()
1853
+ _ = (torch.ones(1, device=dml_device) + 1).item()
1854
+ dml_ok = True
1855
+ except Exception:
1856
+ dml_ok = False
1857
+
1858
+ if cuda_ok:
1859
+ status_cb(f"PyTorch CUDA available: True | device={torch.cuda.get_device_name(0)}")
1860
+ elif mps_ok:
1861
+ status_cb("PyTorch MPS (Apple) available: True")
1862
+ elif dml_ok:
1863
+ status_cb("PyTorch DirectML (Windows) available: True")
1864
+ else:
1865
+ status_cb("PyTorch present, using CPU backend.")
1866
+
1867
+ status_cb(
1868
+ f"PyTorch {getattr(torch, '__version__', '?')} backend: "
1869
+ + ("CUDA" if cuda_ok else "MPS" if mps_ok else "DirectML" if dml_ok else "CPU")
1870
+ )
1871
+ except Exception as e:
1872
+ TORCH_OK = False
1873
+ status_cb(f"PyTorch not available → CPU path. ({e})")
1874
+
1875
+ use_torch = bool(TORCH_OK)
1876
+ if use_torch:
1877
+ # ----- Precision policy (Sport mode but strict FP32) -----
1878
+ try:
1879
+ # Keep autotune for speed
1880
+ torch.backends.cudnn.benchmark = True
1881
+
1882
+ # Force true FP32 everywhere (no TF32 shortcuts)
1883
+ if hasattr(torch.backends, "cudnn"):
1884
+ torch.backends.cudnn.allow_tf32 = False
1885
+ if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
1886
+ torch.backends.cuda.matmul.allow_tf32 = False
1887
+ if hasattr(torch, "set_float32_matmul_precision"):
1888
+ torch.set_float32_matmul_precision("highest")
1889
+ except Exception:
1890
+ pass
1891
+
1892
+ # (optional: telemetry)
1893
+ try:
1894
+ c_tf32 = getattr(torch.backends.cudnn, "allow_tf32", None)
1895
+ m_tf32 = getattr(getattr(torch.backends.cuda, "matmul", object()), "allow_tf32", None)
1896
+ status_cb(
1897
+ f"Precision: cudnn.allow_tf32={c_tf32} | "
1898
+ f"matmul.allow_tf32={m_tf32} | "
1899
+ f"benchmark={torch.backends.cudnn.benchmark}"
1900
+ )
1901
+ except Exception:
1902
+ pass
1903
+
1904
+ _process_gui_events_safely()
1905
+
1906
+ # PSFs (auto-size per frame) + flipped copies
1907
+ psf_out_dir = None
1908
+ psfs, masks_auto, vars_auto = _build_psf_and_assets(
1909
+ paths,
1910
+ make_masks=bool(use_star_masks),
1911
+ make_varmaps=bool(use_variance_maps),
1912
+ status_cb=status_cb,
1913
+ save_dir=None,
1914
+ star_mask_cfg=star_mask_cfg,
1915
+ varmap_cfg=varmap_cfg,
1916
+ star_mask_ref_path=star_mask_ref_path,
1917
+ # NEW:
1918
+ Ht=Ht, Wt=Wt, color_mode=color_mode,
1919
+ )
1920
+
1921
+ # >>> SR: lift PSFs to super-res if requested
1922
+ r = int(max(1, super_res_factor))
1923
+ if r > 1:
1924
+ status_cb(f"MFDeconv: Super-resolution r={r} with σ={sr_sigma} — solving SR PSFs…")
1925
+ _process_gui_events_safely()
1926
+ sr_psfs = []
1927
+ for i, k_native in enumerate(psfs, start=1):
1928
+ h = _solve_super_psf_from_native(k_native, r=r, sigma=float(sr_sigma),
1929
+ iters=int(sr_psf_opt_iters), lr=float(sr_psf_opt_lr))
1930
+ sr_psfs.append(h)
1931
+ status_cb(f" SR-PSF{i}: native {k_native.shape[0]} → {h.shape[0]} (sum={h.sum():.6f})")
1932
+ psfs = sr_psfs
1933
+
1934
+ flip_psf = [_flip_kernel(k) for k in psfs]
1935
+ _emit_pct(0.20, "PSF Ready")
1936
+
1937
+ # Normalize layout BEFORE size harmonization
1938
+ data = _normalize_layout_batch(ys_raw, color_mode) # list of (H,W) or (3,H,W)
1939
+ if str(color_mode).lower() != "luma":
1940
+ # Force strict CHW for every frame
1941
+ data = [_as_chw(a) for a in data]
1942
+ Cs = {a.shape[0] for a in data}
1943
+ if len(Cs) != 1:
1944
+ raise ValueError(f"Inconsistent channel counts in PerChannel mode: {Cs}")
1945
+ _emit_pct(0.25, "Calculating Seed Image...")
1946
+
1947
+ # Center-crop all to common intersection
1948
+ Ht, Wt = _common_hw(data)
1949
+ if any(((a.shape[-2] != Ht) or (a.shape[-1] != Wt)) for a in data):
1950
+ status_cb(f"MFDeconv: Standardizing shapes → crop to {Ht}×{Wt}")
1951
+ data = [_center_crop(a, Ht, Wt) for a in data]
1952
+
1953
+ # Numeric hygiene
1954
+ data = [_sanitize_numeric(a) for a in data]
1955
+
1956
+ # --- SR/native seed ---
1957
+ # --- Seed (choose robust μ-σ or median) ---
1958
+ seed_mode_s = str(seed_mode).lower().strip()
1959
+ if seed_mode_s not in ("robust","median"):
1960
+ seed_mode_s = "robust"
1961
+
1962
+ if seed_mode_s == "median":
1963
+ status_cb("MFDeconv: Building median seed (in-memory)…")
1964
+ # Use already normalized, cropped, sanitized frames
1965
+ seed_native = _seed_median_full_from_data(data)
1966
+ else:
1967
+ status_cb("MFDeconv: Building robust seed (live μ-σ stacking)…")
1968
+ seed_native = _build_seed_running_mu_sigma_from_paths(
1969
+ paths, Ht, Wt, color_mode,
1970
+ bootstrap_frames=20, clip_sigma=5.0,
1971
+ status_cb=status_cb, progress_cb=lambda f,m='': None
1972
+ )
1973
+ if r > 1:
1974
+ if seed_native.ndim == 2:
1975
+ x = _upsample_sum(seed_native / (r*r), r, target_hw=(Ht*r, Wt*r))
1976
+ else:
1977
+ C, Hn, Wn = seed_native.shape
1978
+ x = np.stack(
1979
+ [_upsample_sum(seed_native[c] / (r*r), r, target_hw=(Hn*r, Wn*r)) for c in range(C)],
1980
+ axis=0
1981
+ )
1982
+ else:
1983
+ x = seed_native
1984
+ Hs, Ws = x.shape if x.ndim == 2 else x.shape[-2:]
1985
+
1986
+ # masks/vars aligned to native grid (2D each)
1987
+ auto_masks = masks_auto if use_star_masks else None
1988
+ auto_vars = vars_auto if use_variance_maps else None
1989
+ mask_list = _ensure_mask_list(masks if masks is not None else auto_masks, data)
1990
+ var_list = _ensure_var_list(variances if variances is not None else auto_vars, data)
1991
+
1992
+ iter_dir = None
1993
+ hdr0_seed = None
1994
+ if save_intermediate:
1995
+ iter_dir = _iter_folder(out_path)
1996
+ status_cb(f"MFDeconv: Intermediate outputs → {iter_dir}")
1997
+ try:
1998
+ hdr0_seed = _safe_primary_header(paths[0])
1999
+ except Exception:
2000
+ hdr0_seed = fits.Header()
2001
+ _save_iter_image(x, hdr0_seed, iter_dir, "seed", color_mode)
2002
+
2003
+ status_cb("MFDeconv: Calculating Backgrounds and MADs…")
2004
+ _process_gui_events_safely()
2005
+ bg_est = _sep_bg_rms(data) or (np.median([np.median(np.abs(y - np.median(y))) for y in (data if isinstance(data, list) else [data])]) * 1.4826)
2006
+ status_cb(f"MFDeconv: color_mode={color_mode}, huber_delta={huber_delta} (bg RMS~{bg_est:.3g})")
2007
+ _process_gui_events_safely()
2008
+
2009
+ status_cb("Computing FFTs and Allocating Scratch…")
2010
+ _process_gui_events_safely()
2011
+
2012
+ # -------- precompute and allocate scratch --------
2013
+ per_frame_logging = (r > 1)
2014
+ if use_torch:
2015
+ x_t = _to_t(_contig(x))
2016
+ num = torch.zeros_like(x_t)
2017
+ den = torch.zeros_like(x_t)
2018
+
2019
+ if r > 1:
2020
+ # >>> SR path now uses SPATIAL CONV (cuDNN) to avoid huge FFT buffers
2021
+ psf_t = [_to_t(_contig(k))[None, None] for k in psfs] # (1,1,kh,kw)
2022
+ psfT_t = [_to_t(_contig(kT))[None, None] for kT in flip_psf]
2023
+ else:
2024
+ # Native spatial (as before)
2025
+ psf_t = [_to_t(_contig(k))[None, None] for k in psfs]
2026
+ psfT_t = [_to_t(_contig(kT))[None, None] for kT in flip_psf]
2027
+
2028
+ else:
2029
+ x_t = x
2030
+ # CPU path: keep your (more RAM-tolerant) FFT packs
2031
+ if r > 1:
2032
+ if x_t.ndim == 2:
2033
+ Hs, Ws = x_t.shape
2034
+ else:
2035
+ _, Hs, Ws = x_t.shape
2036
+ Kfs, KTfs, meta = _precompute_np_psf_ffts(psfs, flip_psf, Hs, Ws)
2037
+ pred_super = np.empty_like(x_t)
2038
+ tmp_out = np.empty_like(x_t)
2039
+ else:
2040
+ Kfs, KTfs, meta = _precompute_np_psf_ffts(psfs, flip_psf, Hs, Ws)
2041
+ pred_super = np.empty_like(x_t)
2042
+ tmp_out = np.empty_like(x_t)
2043
+ num = np.zeros_like(x_t)
2044
+ den = np.zeros_like(x_t)
2045
+
2046
+ status_cb("Starting First Multiplicative Iteration…")
2047
+ _process_gui_events_safely()
2048
+
2049
+ cm = _safe_inference_context() if use_torch else NO_GRAD
2050
+ rho_is_l2 = (str(rho).lower() == "l2")
2051
+ local_delta = 0.0 if rho_is_l2 else huber_delta
2052
+ used_iters = 0
2053
+ early_stopped = False
2054
+
2055
+ auto_delta_cache = None
2056
+ if use_torch and (huber_delta < 0) and (not rho_is_l2):
2057
+ auto_delta_cache = [None] * len(paths)
2058
+
2059
+ early = EarlyStopper(
2060
+ tol_upd_floor=2e-4, # match new version
2061
+ tol_rel_floor=5e-4,
2062
+ early_frac=0.40,
2063
+ ema_alpha=0.5,
2064
+ patience=2,
2065
+ min_iters=min_iters,
2066
+ status_cb=status_cb
2067
+ )
2068
+ x_ndim = 2 if (np.ndim(x) == 2) else 3
2069
+ fixed = 0
2070
+ for i, a in enumerate(data):
2071
+ if a.ndim != x_ndim:
2072
+ # fix common mono cases only
2073
+ if x_ndim == 2 and a.ndim == 3 and a.shape[0] == 1:
2074
+ data[i] = a[0]; fixed += 1
2075
+ elif x_ndim == 2 and a.ndim == 3 and a.shape[-1] == 1:
2076
+ data[i] = a[..., 0]; fixed += 1
2077
+
2078
+ with cm():
2079
+ for it in range(1, max_iters + 1):
2080
+ if use_torch:
2081
+ num.zero_(); den.zero_()
2082
+
2083
+ if r > 1:
2084
+ # -------- SR path (SPATIAL conv + stream) --------
2085
+ for fidx, (wk, wkT) in enumerate(zip(psf_t, psfT_t)):
2086
+ yt_np = data[fidx] # CHW or HW (CPU)
2087
+ mt_np = mask_list[fidx]
2088
+ vt_np = var_list[fidx]
2089
+
2090
+ yt = torch.as_tensor(yt_np, dtype=x_t.dtype, device=x_t.device)
2091
+ mt = None if mt_np is None else torch.as_tensor(mt_np, dtype=x_t.dtype, device=x_t.device)
2092
+ vt = None if vt_np is None else torch.as_tensor(vt_np, dtype=x_t.dtype, device=x_t.device)
2093
+
2094
+ # SR conv on grid of x_t
2095
+ pred_super = _conv_same_torch(x_t, wk) # SR grid
2096
+ pred_low = _downsample_avg_t(pred_super, r) # native grid
2097
+
2098
+ if auto_delta_cache is not None:
2099
+ if (auto_delta_cache[fidx] is None) or (it % 5 == 1):
2100
+ rnat = yt - pred_low
2101
+ med = torch.median(rnat)
2102
+ mad = torch.median(torch.abs(rnat - med)) + 1e-6
2103
+ rms = 1.4826 * mad
2104
+ auto_delta_cache[fidx] = float((-huber_delta) * torch.clamp(rms, min=1e-6).item())
2105
+ wmap_low = _weight_map(yt, pred_low, auto_delta_cache[fidx], var_map=vt, mask=mt)
2106
+ else:
2107
+ wmap_low = _weight_map(yt, pred_low, local_delta, var_map=vt, mask=mt)
2108
+
2109
+ # lift back to SR via sum-replicate
2110
+ up_y = _upsample_sum_t(wmap_low * yt, r)
2111
+ up_pred = _upsample_sum_t(wmap_low * pred_low, r)
2112
+
2113
+ # accumulate via adjoint kernel on SR grid
2114
+ num += _conv_same_torch(up_y, wkT)
2115
+ den += _conv_same_torch(up_pred, wkT)
2116
+
2117
+ # free temps as aggressively as possible
2118
+ del yt, mt, vt, pred_super, pred_low, wmap_low, up_y, up_pred
2119
+ if cuda_ok:
2120
+ try: torch.cuda.empty_cache()
2121
+ except Exception as e:
2122
+ import logging
2123
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
2124
+
2125
+ if per_frame_logging and ((fidx & 7) == 0):
2126
+ status_cb(f"Iter {it}/{max_iters} — frame {fidx+1}/{len(paths)} (SR spatial)")
2127
+
2128
+ else:
2129
+ # -------- Native path (spatial conv, stream) --------
2130
+ for fidx, (wk, wkT) in enumerate(zip(psf_t, psfT_t)):
2131
+ yt_np = data[fidx]
2132
+ mt_np = mask_list[fidx]
2133
+ vt_np = var_list[fidx]
2134
+
2135
+ yt = torch.as_tensor(yt_np, dtype=x_t.dtype, device=x_t.device)
2136
+ mt = None if mt_np is None else torch.as_tensor(mt_np, dtype=x_t.dtype, device=x_t.device)
2137
+ vt = None if vt_np is None else torch.as_tensor(vt_np, dtype=x_t.dtype, device=x_t.device)
2138
+
2139
+ pred = _conv_same_torch(x_t, wk)
2140
+ wmap = _weight_map(yt, pred, local_delta, var_map=vt, mask=mt)
2141
+ up_y = wmap * yt
2142
+ up_pred = wmap * pred
2143
+ num += _conv_same_torch(up_y, wkT)
2144
+ den += _conv_same_torch(up_pred, wkT)
2145
+
2146
+ del yt, mt, vt, pred, wmap, up_y, up_pred
2147
+ if cuda_ok:
2148
+ try: torch.cuda.empty_cache()
2149
+ except Exception as e:
2150
+ import logging
2151
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
2152
+
2153
+ ratio = num / (den + EPS)
2154
+ neutral = (den.abs() < 1e-12) & (num.abs() < 1e-12)
2155
+ ratio = torch.where(neutral, torch.ones_like(ratio), ratio)
2156
+ upd = torch.clamp(ratio, 1.0 / kappa, kappa)
2157
+ x_next = torch.clamp(x_t * upd, min=0.0)
2158
+
2159
+ upd_med = torch.median(torch.abs(upd - 1))
2160
+ rel_change = (torch.median(torch.abs(x_next - x_t)) /
2161
+ (torch.median(torch.abs(x_t)) + 1e-8))
2162
+
2163
+ # candidates for convergence
2164
+ try:
2165
+ um = float(upd_med.detach().cpu().item())
2166
+ except Exception:
2167
+ um = float(upd_med)
2168
+
2169
+ try:
2170
+ rc = float(rel_change.detach().cpu().item())
2171
+ except Exception:
2172
+ rc = float(rel_change)
2173
+
2174
+ if early.step(it, max_iters, um, rc):
2175
+ x_t = x_next
2176
+ used_iters = it
2177
+ early_stopped = True
2178
+ _process_gui_events_safely()
2179
+ break
2180
+
2181
+ x_t = (1.0 - relax) * x_t + relax * x_next
2182
+
2183
+ else:
2184
+ # -------- NumPy path (unchanged) --------
2185
+ num.fill(0.0); den.fill(0.0)
2186
+ if r > 1:
2187
+ for (Kf, KTf, (kh, kw, fftH, fftW)), m2d, v2d, y_nat in zip(
2188
+ zip(Kfs, KTfs, meta), mask_list, var_list, data):
2189
+ _fft_conv_same_np(x_t, Kf, kh, kw, fftH, fftW, pred_super)
2190
+ pred_low = _downsample_avg(pred_super, r)
2191
+ wmap_low = _weight_map(y_nat, pred_low, local_delta, var_map=v2d, mask=m2d)
2192
+ up_y = _upsample_sum(wmap_low * y_nat, r, target_hw=pred_super.shape[-2:])
2193
+ up_pred = _upsample_sum(wmap_low * pred_low, r, target_hw=pred_super.shape[-2:])
2194
+ _fft_conv_same_np(up_y, KTf, kh, kw, fftH, fftW, tmp_out); num += tmp_out
2195
+ _fft_conv_same_np(up_pred, KTf, kh, kw, fftH, fftW, tmp_out); den += tmp_out
2196
+ else:
2197
+ for (Kf, KTf, (kh, kw, fftH, fftW)), m2d, v2d, y_nat in zip(
2198
+ zip(Kfs, KTfs, meta), mask_list, var_list, data):
2199
+ _fft_conv_same_np(x_t, Kf, kh, kw, fftH, fftW, pred_super)
2200
+ pred = pred_super
2201
+ wmap = _weight_map(y_nat, pred, local_delta, var_map=v2d, mask=m2d)
2202
+ up_y, up_pred = (wmap * y_nat), (wmap * pred)
2203
+ _fft_conv_same_np(up_y, KTf, kh, kw, fftH, fftW, tmp_out); num += tmp_out
2204
+ _fft_conv_same_np(up_pred, KTf, kh, kw, fftH, fftW, tmp_out); den += tmp_out
2205
+
2206
+ ratio = num / (den + EPS)
2207
+ neutral = (np.abs(den) < 1e-12) & (np.abs(num) < 1e-12)
2208
+ ratio[neutral] = 1.0
2209
+ upd = np.clip(ratio, 1.0 / kappa, kappa)
2210
+ x_next = np.clip(x_t * upd, 0.0, None)
2211
+
2212
+ upd_med = np.median(np.abs(upd - 1.0))
2213
+ rel_change = (np.median(np.abs(x_next - x_t)) /
2214
+ (np.median(np.abs(x_t)) + 1e-8))
2215
+ um = float(upd_med)
2216
+ rc = float(rel_change)
2217
+
2218
+ if early.step(it, max_iters, um, rc):
2219
+ x_t = x_next
2220
+ used_iters = it
2221
+ early_stopped = True
2222
+ _process_gui_events_safely()
2223
+ break
2224
+
2225
+
2226
+ x_t = (1.0 - relax) * x_t + relax * x_next
2227
+
2228
+ # save intermediate
2229
+ if save_intermediate and (it % int(max(1, save_every)) == 0):
2230
+ try:
2231
+ x_np = x_t.detach().cpu().numpy().astype(np.float32) if use_torch else x_t.astype(np.float32)
2232
+ _save_iter_image(x_np, hdr0_seed, iter_dir, f"iter_{it:03d}", color_mode)
2233
+ except Exception as _e:
2234
+ status_cb(f"Intermediate save failed at iter {it}: {_e}")
2235
+
2236
+ frac = 0.25 + 0.70 * (it / float(max_iters))
2237
+ _emit_pct(frac, f"Iteration {it}/{max_iters}")
2238
+
2239
+ _process_gui_events_safely()
2240
+
2241
+ if not early_stopped:
2242
+ used_iters = max_iters
2243
+
2244
+ # ----------------------------
2245
+ # Save result (keep FITS-friendly order: (C,H,W))
2246
+ # ----------------------------
2247
+ _emit_pct(0.97, "saving")
2248
+ x_final = x_t.detach().cpu().numpy().astype(np.float32) if use_torch else x_t.astype(np.float32)
2249
+
2250
+ if x_final.ndim == 3:
2251
+ if x_final.shape[0] not in (1, 3) and x_final.shape[-1] in (1, 3):
2252
+ x_final = np.moveaxis(x_final, -1, 0)
2253
+ if x_final.shape[0] == 1:
2254
+ x_final = x_final[0]
2255
+
2256
+ try:
2257
+ hdr0 = _safe_primary_header(paths[0])
2258
+ except Exception:
2259
+ hdr0 = fits.Header()
2260
+ hdr0['MFDECONV'] = (True, 'Seti Astro multi-frame deconvolution')
2261
+ hdr0['MF_COLOR'] = (str(color_mode), 'Color mode used')
2262
+ hdr0['MF_RHO'] = (str(rho), 'Loss: huber|l2')
2263
+ hdr0['MF_HDEL'] = (float(huber_delta), 'Huber delta (>0 abs, <0 autoxRMS)')
2264
+ hdr0['MF_MASK'] = (bool(use_star_masks), 'Used auto star masks')
2265
+ hdr0['MF_VAR'] = (bool(use_variance_maps), 'Used auto variance maps')
2266
+
2267
+ hdr0['MF_SR'] = (int(r), 'Super-resolution factor (1 := native)')
2268
+ if r > 1:
2269
+ hdr0['MF_SRSIG'] = (float(sr_sigma), 'Gaussian sigma for SR PSF fit (pixels at native)')
2270
+ hdr0['MF_SRIT'] = (int(sr_psf_opt_iters), 'SR-PSF solver iters')
2271
+
2272
+ hdr0['MF_ITMAX'] = (int(max_iters), 'Requested max iterations')
2273
+ hdr0['MF_ITERS'] = (int(used_iters), 'Actual iterations run')
2274
+ hdr0['MF_ESTOP'] = (bool(early_stopped), 'Early stop triggered')
2275
+
2276
+ if isinstance(x_final, np.ndarray):
2277
+ if x_final.ndim == 2:
2278
+ hdr0['MF_SHAPE'] = (f"{x_final.shape[0]}x{x_final.shape[1]}", 'Saved as 2D image (HxW)')
2279
+ elif x_final.ndim == 3:
2280
+ C, H, W = x_final.shape
2281
+ hdr0['MF_SHAPE'] = (f"{C}x{H}x{W}", 'Saved as 3D cube (CxHxW)')
2282
+
2283
+ save_path = _sr_out_path(out_path, super_res_factor)
2284
+ safe_out_path = _nonclobber_path(str(save_path))
2285
+ if safe_out_path != str(save_path):
2286
+ status_cb(f"Output exists — saving as: {safe_out_path}")
2287
+ fits.PrimaryHDU(data=x_final, header=hdr0).writeto(safe_out_path, overwrite=False)
2288
+
2289
+ status_cb(f"✅ MFDeconv saved: {safe_out_path} (iters used: {used_iters}{', early stop' if early_stopped else ''})")
2290
+ _emit_pct(1.00, "done")
2291
+ _process_gui_events_safely()
2292
+
2293
+ try:
2294
+ if use_torch:
2295
+ try: del num, den
2296
+ except Exception as e:
2297
+ import logging
2298
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
2299
+ try: del psf_t, psfT_t
2300
+ except Exception as e:
2301
+ import logging
2302
+ logging.debug(f"Exception suppressed: {type(e).__name__}: {e}")
2303
+ _free_torch_memory()
2304
+ except Exception:
2305
+ pass
2306
+
2307
+ return safe_out_path
2308
+
2309
+
2310
+
2311
+ # -----------------------------
2312
+ # Worker
2313
+ # -----------------------------
2314
+
2315
+ class MultiFrameDeconvWorkerSport(QObject):
2316
+ progress = pyqtSignal(str)
2317
+ finished = pyqtSignal(bool, str, str) # success, message, out_path
2318
+
2319
+ def __init__(self, parent, aligned_paths, output_path, iters, kappa, color_mode,
2320
+ huber_delta, min_iters, use_star_masks=False, use_variance_maps=False, rho="huber",
2321
+ star_mask_cfg: dict | None = None, varmap_cfg: dict | None = None,
2322
+ save_intermediate: bool = False,
2323
+ seed_mode: str = "robust",
2324
+ # NEW SR params
2325
+ super_res_factor: int = 1,
2326
+ sr_sigma: float = 1.1,
2327
+ sr_psf_opt_iters: int = 250,
2328
+ sr_psf_opt_lr: float = 0.1,
2329
+ star_mask_ref_path: str | None = None):
2330
+ super().__init__(parent)
2331
+ self.aligned_paths = aligned_paths
2332
+ self.output_path = output_path
2333
+ self.iters = iters
2334
+ self.kappa = kappa
2335
+ self.color_mode = color_mode
2336
+ self.huber_delta = huber_delta
2337
+ self.min_iters = min_iters # NEW
2338
+ self.star_mask_cfg = star_mask_cfg or {}
2339
+ self.varmap_cfg = varmap_cfg or {}
2340
+ self.use_star_masks = use_star_masks
2341
+ self.use_variance_maps = use_variance_maps
2342
+ self.rho = rho
2343
+ self.save_intermediate = save_intermediate
2344
+ self.super_res_factor = int(super_res_factor)
2345
+ self.sr_sigma = float(sr_sigma)
2346
+ self.sr_psf_opt_iters = int(sr_psf_opt_iters)
2347
+ self.sr_psf_opt_lr = float(sr_psf_opt_lr)
2348
+ self.star_mask_ref_path = star_mask_ref_path
2349
+ self.seed_mode = seed_mode
2350
+
2351
+
2352
+ def _log(self, s): self.progress.emit(s)
2353
+
2354
+ def run(self):
2355
+ try:
2356
+ out = multiframe_deconv(
2357
+ self.aligned_paths,
2358
+ self.output_path,
2359
+ iters=self.iters,
2360
+ kappa=self.kappa,
2361
+ color_mode=self.color_mode,
2362
+ seed_mode=self.seed_mode,
2363
+ huber_delta=self.huber_delta,
2364
+ use_star_masks=self.use_star_masks,
2365
+ use_variance_maps=self.use_variance_maps,
2366
+ rho=self.rho,
2367
+ min_iters=self.min_iters,
2368
+ status_cb=self._log,
2369
+ star_mask_cfg=self.star_mask_cfg,
2370
+ varmap_cfg=self.varmap_cfg,
2371
+ save_intermediate=self.save_intermediate,
2372
+ # NEW SR forwards
2373
+ super_res_factor=self.super_res_factor,
2374
+ sr_sigma=self.sr_sigma,
2375
+ sr_psf_opt_iters=self.sr_psf_opt_iters,
2376
+ sr_psf_opt_lr=self.sr_psf_opt_lr,
2377
+ star_mask_ref_path=self.star_mask_ref_path,
2378
+ )
2379
+ self.finished.emit(True, "MF deconvolution complete.", out)
2380
+ _process_gui_events_safely()
2381
+ except Exception as e:
2382
+ self.finished.emit(False, f"MF deconvolution failed: {e}", "")