R.A.'s Aberration Correction (AI)

AI-powered correction of optical aberrations. Runs Riccardo Alberghi's ONNX model over your image to clean up coma, astigmatism, field curvature artifacts, and other distortions that leave stars elongated, tailed, or seagulled toward the frame edges. SASpro downloads the latest published model automatically, tiles the image with Hann-windowed overlap for seamless blending, picks the best ONNX execution provider for your hardware (DirectML / CUDA / CoreML / CPU), and preserves a thin ring of original pixels at the border where the model can't see enough context. Supports custom model files, mask-blended application, and preset drag-and-drop to replay the same settings against any image.
Model and training by Riccardo Alberghi (github.com/riccardoalberghi). SASpro integrates his published ONNX models through a hardened runner that handles automatic download, cross-platform execution-provider selection, image tiling for inputs larger than the model's native tile size, and a durable self-patch for an ONNX opset-10 incompatibility that would otherwise prevent loading on modern onnxruntime. The AI engine is Riccardo's work; SASpro's job is to make it reliable, fast, and usable from inside your workflow.

1. The core idea

Optical systems — refractors, Newtonians, SCTs with reducers, everything — produce field-dependent aberrations. Stars near the frame center are tight circles; stars near the corners stretch into coma tails, seagulls, or triangles depending on the system. Classical corrections (field flatteners, coma correctors) help but rarely eliminate the effect, especially with budget optics or tight integration time. Aberration AI learns the mapping from distorted field positions to corrected ones from a training set of real astrophotos, and runs that correction directly on your image.

The model is a convolutional network that processes fixed-size patches (usually 512×512) and outputs the same-size corrected patch. SASpro's job is to feed the full image through in overlapping tiles, blend them smoothly, and bolt on the pre- and post-processing (normalization, border preservation, mask blending, provider selection) that turn a research-grade ONNX file into a reliable processing step.

2. First run — getting the model

Three ways to end up with a working model:

The Browse button is smart. If you pick a file that lives inside the SASpro models folder, it's registered as the "downloaded" active model. If you pick a file anywhere else, it's registered as a "custom" model and the custom checkbox flips on automatically. One button, two roles — SASpro figures out which from the file's location.

3. Provider selection

ONNX models can execute on different hardware backends. SASpro picks sensibly per platform:

The Auto GPU (if available) checkbox picks whatever's best available. Uncheck it and use the Provider dropdown to force a specific backend — useful for debugging ("does this still fail on CPU?") or when you want to leave the GPU free for something else.

CUDA fallback mid-run. If a CUDA session loads successfully but then fails during inference with a kernel-image or arch mismatch (common when your onnxruntime-gpu was built for a different CUDA toolkit than your driver), SASpro catches that error, logs it, and restarts the run on CPU automatically. You don't lose your in-progress work; the whole image reprocesses on CPU and the result lands as usual. Watch the status log for "CUDA kernel mismatch detected — falling back to CPU…".

4. Patch and overlap

Models are trained on a specific input size (currently 512×512). SASpro inspects the loaded model via the ONNX graph and automatically snaps the Patch spinbox to the required value if the model has a fixed input shape — you can't accidentally run a 512-trained model with 256 tiles and get degraded output.

Overlap controls how much adjacent tiles share. The default 64 pixels at 512-tile size means each tile contributes to the final output across ~12% of its area as a weighted blend with its neighbors. Lower overlap = fewer tiles, faster run, but visible seams at tile boundaries. Higher overlap = smoother blend, slower run, more memory pressure. For large images on slow hardware, bump overlap to 128 if you can see seams; drop to 32 if you're desperate for speed on obvious aberrations where seams won't matter.

5. The CoreML tile-size guard

CoreML on Intel macOS has a hard practical limit on tile size — it can execute the model, but memory allocation inside the CoreML runtime starts failing above ~128 px per side. SASpro checks the model's required patch size against this limit: if the model needs 512-pixel tiles but CoreML is the only non-CPU option, SASpro falls back to CPU and logs "CoreML limited to small tiles; model requires 512px → using CPU." No crash, no partial output, just the right backend for the job.

6. Border preservation

After the full tiled run, SASpro copies a 10-pixel ring from the original image back over the output. This fixes the one place where the network can't see enough context: the first few pixels at the frame edge, where the model's receptive field doesn't have real pixels to feed its convolutions and tends to produce dim or washed-out borders. The ring is thin enough that you'll never notice the transition in practice; it's wide enough to swallow the artifact completely.

If the output shape doesn't match the input (shouldn't happen, but defensively) the border copy is silently skipped.

7. Masks and selective correction

If the active document has a mask active, SASpro blends the AI output with the original using the mask as alpha:

\mathrm{out} = m \cdot \mathrm{processed} + (1 - m) \cdot \mathrm{src}

Mask values are normalized to [0, 1] (handles uint masks and float masks of any range), resampled to the output size if needed, and broadcast across channels for RGB inputs. White areas of the mask get full AI correction; black areas keep their original pixels; grays blend smoothly. Perfect for "I only want to clean up the corners, not touch the galaxy core" — paint a mask that excludes the center and apply.

8. Mask-aware application history

The result is applied through SASpro's normal undo-aware document pipeline (step name: "Aberration AI"), so Ctrl+Z gets you back to the original. The processing parameters block records the exact model path, patch size, overlap, provider, border width, and mask ID used for the run — so if you load the project months later you can see precisely what was done.

9. The preset drag handle

The small grip icon at the bottom-left of the dialog is a draggable preset handle. Drag it onto the SASpro canvas to create a shortcut button pre-wired to the Aberration AI tool with your current settings baked in. Drop it directly onto an image to run the current preset against that image headlessly — no dialog, just processing. The preset captures: model path, patch size, overlap, border width, Auto GPU setting, and (if Auto GPU is off) the specific provider.

SASpro also registers the last run as the "Replay Last Action" command, so hitting that shortcut from anywhere in SASpro re-runs Aberration AI with your most recent settings against the current view.

10. The automatic ONNX TopK patch (self-healing)

This is a surprise technical highlight. Some published Aberration AI models were exported with ONNX opset 10 (or earlier), when the TopK operator still allowed its K input as a 0-dimensional scalar tensor. The ONNX specification was tightened in opset 11 to require K as a 1-dimensional tensor of size 1, and modern onnxruntime (version 1.14 and later) enforces this at graph-validation time — refusing to load the model at all, before even checking which execution provider you asked for. The failure looks the same on GPU and CPU because it happens during the parse step, not inference.

Since the model is published by a third party, SASpro can't just re-export it. Instead, on the first failed load, SASpro's _patch_topk_k_shape opens the model file, finds every TopK node, rewrites each K input to a shape-[1] tensor, saves the patched file back to disk, verifies the patch actually landed, and retries the load. The fix is durable — subsequent launches load without touching the self-heal code at all.

The patch handles every ONNX-legal way a scalar K can be encoded: initializer tensors (0-D), Constant nodes with value / value_int / value_ints / value_float attributes, and runtime-computed K from Shape/Gather chains (which can't be rewritten in place, so SASpro inserts a Reshape([1]) node between the producer and the TopK instead). You'll see "✅ patched N TopK K input(s)" in the log on the first run after installing a new model; after that, silent success.

11. Workflow recipes

First-light on a new scope. Load your stacked integration. Open Aberration AI, click Run. If no model is present, SASpro downloads it and processes in one go. Compare the corner stars before and after — if the correction looks right, push to the view and continue your normal workflow. If it looks over- or under-corrected, blend with a mask (next recipe).
Partial correction via mask. Create a mask that's white at the corners and transitions to black at the center (a radial gradient from a Range Mask tool, inverted). Enable the mask on the active document. Open Aberration AI, Run. Only the corners get the AI correction; the center stays as your original data. Perfect for scopes where the center is already tight and you only need help at the edges.
Side-by-side provider comparison. Run once with Auto GPU on, note the time logged in the status bar. Uncheck Auto GPU, pick "CPUExecutionProvider" from the dropdown, run again. Compare timings. On a GPU-equipped Windows or Linux box, DirectML / CUDA will be 5-20× faster than CPU for the same output. Useful for deciding whether to buy a GPU for this specific workflow.
Custom model experiment. Download a community-trained variant (someone's model for a specific scope or sensor). Browse to it. The checkbox flips on automatically. Run. Compare to the stock model's output. Toggle the checkbox back off to A/B against the original any time without re-selecting files.
Batch across a session. Open each night's stack, run Aberration AI once with the settings you like. On the second image, drag the preset grip from the bottom-left of the dialog onto the view — the correction runs headlessly with those exact settings. No dialog, no clicks beyond the drop. Or set up a shortcut button on the SASpro canvas by dropping the grip there, and click once per image.
Replay last. After the first Aberration AI run of your session, every subsequent image can get the same correction via SASpro's Replay Last Action command (your bound shortcut). SASpro remembers the model, patch, overlap, border, and provider from the last run and applies them in one keystroke.

12. Credits

Model and training. Riccardo Alberghi — github.com/riccardoalberghi. The ONNX model, its architecture, training data, and all weights are Riccardo's work. SASpro's integration is a thin runner around his published artifacts.

SASpro integration. Written by Franklin Marek as part of SetiAstro Suite. Includes: automatic model download from GitHub Releases, cross-platform ONNX execution-provider selection (DirectML / CUDA / CoreML / CPU), Hann-windowed tiled inference with mono/RGB handling, uint16 / float32 normalization, border preservation, mask-aware blending, undo-aware application via SASpro's document pipeline, preset drag-and-drop, Replay Last Action registration, and the automatic TopK opset-10 compatibility patch described in §10.

ONNX Runtime — onnxruntime.ai. Execution providers: DirectML (Windows, DX12 GPUs), CUDA (NVIDIA GPUs on Windows/Linux), CoreML (macOS), CPU (universal). The onnx Python package handles the self-heal rewrite of the model graph.

Techniques. Hann windowing follows the standard overlap-add synthesis for 2D image reconstruction; the implementation is pure NumPy. ONNX graph surgery for runtime-computed K inserts a Reshape node from the ONNX standard operator set. GitHub Releases API consumed via the requests package with streamed download and configurable chunk size.

Technical details

13.1. Hann window blending

Each tile is multiplied by a 2D Hann (also called Hanning) window before accumulation. The 1D Hann window for N samples:

w(n) = 0.5 \cdot (1 - \cos(2\pi n / (N - 1)))

Peaks at 0.5 in the middle and tapers to 0 at both ends. The 2D version is the outer product:

w_{2D}(i, j) = w(i) \cdot w(j)

Peaks at 0.25 in the tile center and tapers to 0 at all four edges and corners. This is the standard overlap-add weighting for seamless image reconstruction — adjacent tiles' windows sum to a near-constant across the overlap region, so the weighted mean below produces a smooth result with no visible seams at tile boundaries.

13.2. Tile layout

For an image dimension N (height or width) and a patch size P with overlap O:

\mathrm{stride} = P - O
\{0,\, \mathrm{stride},\, 2 \cdot \mathrm{stride},\, \ldots\} \cup \{N - P\}

Positions start at 0 and advance by stride, until the next position would push the tile past the image boundary — at which point SASpro appends a final tile anchored at N - P (so the last tile always ends exactly at the image edge). Duplicates (common when N is just slightly larger than a multiple of stride) are removed and the result sorted. The outer product with the per-axis positions gives the full 2D tile grid.

13.3. Weighted-mean reconstruction

Each tile's output is accumulated into the output buffer, multiplied by its Hann window; the window values themselves are accumulated into a weight buffer. At the end, the pixelwise ratio gives the final result:

\mathrm{out}[i, j] = \frac{\sum_k \mathrm{tile}_k[i, j] \cdot w_{2D}[i_k, j_k]}{\sum_k w_{2D}[i_k, j_k]}

Where tile_k is the k-th tile's model output at position (i_k, j_k) within that tile. Weight buffer values that happen to be zero (shouldn't occur with Hann and correct overlap, but defensively) are replaced with 1 to prevent division by zero. The result is the equivalent of a Hann-windowed overlap-add synthesis — standard practice in audio spectral processing, adapted here to 2D image patches.

13.4. Normalization and padding

The model expects float32 inputs in [0, 1]. SASpro handles conversion:

x_{\mathrm{float}} = x_{\mathrm{uint16}} / 65535

and the inverse on output:

x_{\mathrm{uint16}} = \lfloor x_{\mathrm{float}} \cdot 65535 \rfloor

Float inputs pass through directly (clipped to [0, 1] on output). Monochrome images are promoted to a (1, H, W) tensor for processing and demoted back to (H, W) at the end. RGB images are channel-separated: the network processes each channel independently with the same tile schedule, which keeps memory per inference call low and matches how the model was trained (on luminance-equivalent patches). Images smaller than the patch size are edge-padded:

\mathrm{pad}_H = \max(0, P - H)

with symmetric pad in W. The padding is cropped off at the end before the result is returned, so the output is exactly the input size.

13.5. Patch tensor shape

Each tile is reshaped into the batch-of-1, 1-channel 4D form the model expects:

\mathrm{inp} \in \mathbb{R}^{1 \times 1 \times P \times P}

Contiguous memory layout is forced via np.ascontiguousarray before the ONNX call to avoid an internal copy inside the runtime. The output is np.squeeze-d back to (1, P, P) for accumulation.

13.6. Progress tracking

Progress is reported as a fraction after each tile finishes:

\mathrm{frac} = (\mathrm{tiles\_done}) / (|H| \cdot |W| \cdot C)

where |H| and |W| are the number of tile positions along each axis and C is the channel count. For a 512-tile schedule with 64 overlap on an 8K monochrome image, that's ~16×16×1 = 256 tiles; the progress bar updates smoothly across the run. The automatic download path reports its own progress on the same bar, remapped to [0, 0.15] so the AI inference occupies [0.15, 1.0] when both happen.

13.7. Border preservation and mask blend

After the Hann accumulation and clip-to-[0,1], a 10-pixel ring from the original image overwrites the output edge:

\mathrm{dst}[\mathrm{border}] = \mathrm{src}[\mathrm{border}]

Border width is clamped to half the smaller image dimension so it can't exceed the image. If a mask is active, the final mask-blend step applies:

\mathrm{out} = m \cdot \mathrm{processed} + (1 - m) \cdot \mathrm{src}

Mask values are normalized to [0, 1] first (dividing by the dtype max for uint masks, or by the empirical max for floats above 1). Mask resolution is reconciled to the output size by nearest-neighbor subsampling if different. For RGB outputs the mask is broadcast across channels (mono mask applies to all three equally). If the mask is a different shape dimensionality than the output (mono mask + RGB output), the input is dimension-matched via channel repeat or extraction before the blend.

13.8. NaN safety and output clipping

The final step before returning is a combined NaN replacement and range clip:

\mathrm{out} = \mathrm{clip}(\mathrm{nan\_to\_num}(x),\; 0,\; 1)

Models that occasionally produce NaN values (division instabilities, overflow in intermediate activations) still produce usable output because NaN is replaced with 0 before clipping. Infinities are similarly handled by np.nan_to_num — +inf maps to a large finite value and then clips to 1.

13.9. The TopK auto-patch — details

The ONNX specification for TopK at opset 10 permitted K as a scalar; opset 11 changed this to require a 1-D tensor of size 1. onnxruntime 1.14 and later enforce the current spec during graph validation, before any provider is engaged — the error is identical on CPU and GPU. SASpro detects this error by matching on the common substrings ("TopK" plus one of "one-dimensional", "K input", or "ShapeInferenceError", since different ORT versions phrase it differently).

On detection, _patch_topk_k_shape loads the model with the onnx Python package and walks every TopK node, handling five cases for the K source:

  1. Initializer tensor (0-D) — rewrite the initializer in place as np.array([int(scalar)], dtype=np.int64).
  2. Constant node with value — the attribute carries a TensorProto; if 0-D, rewrite to shape [1] int64.
  3. Constant node with value_int — implicit scalar int64, re-emit as a shape-[1] tensor.
  4. Constant node with value_ints — length 0 or 1 attribute vector, re-emit as a shape-[1] tensor.
  5. Constant node with value_float — implicit scalar float, cast to int64 and re-emit as shape-[1].

Any remaining TopK K inputs that come from runtime computation (Shape → Gather chains, common in PyTorch exports where k depends on an input dimension) can't be rewritten at the source. SASpro inserts a Reshape node between the producer and the TopK:

producer → Reshape(shape=[1]) → TopK.K

which is a no-op if the input was already shape [1] and the required fix if it was rank 0. A shared int64 [1] initializer named _saspro_topk_k_reshape_shape_ is added once; each new Reshape gets a unique output name so multiple affected TopK nodes don't collide.

After the save, SASpro re-loads the file and verifies the patch actually landed on disk — every initializer and Constant feeding a TopK K input must now have dims == [1]. If verification fails, the error is surfaced to the user with both the original ORT error and the patch diagnostic so they can see exactly what was tried. Successful patches are logged with the count of fixed inputs and the file name, and the retry proceeds normally.

13.10. Provider selection logic

The pick_providers function returns an ordered list that ORT tries in sequence, falling back to the next on failure:

Session open uses a try/except around _open_session_with_topk_fix: if the preferred providers fail (missing CUDA DLL, CoreML permission error, etc.), SASpro retries with just [CPUExecutionProvider] and logs the fallback reason. Mid-inference CUDA kernel errors ("cuda" or "kernel image" in the message) trigger a full rerun on CPU — the session is replaced and the tiled inference restarts from tile 0 with the new session.

13.11. QSettings keys used

Three persisted settings control model selection:

Preset drag-and-drop payloads use the resolved model path (whichever is currently active) at the moment of drag, so a shortcut created with the custom model enabled captures the custom path; a shortcut created with the downloaded model active captures that path. Headless consumers of the preset can fall back to resolve_aberration_model_path if the embedded path is empty — useful for shortcuts that should always use "whatever model is currently selected" rather than freezing to a specific file.