SIQ: Super-Resolution Image Quantification

SIQ is a powerful and flexible Python library for deep learning-based image super-resolution, with a focus on medical imaging applications. It provides a comprehensive toolkit for every stage of the super-resolution workflow, from data generation and model training to robust inference and evaluation.

The library is built on TensorFlow/Keras and ANTSpy and is designed to handle complex, real-world challenges such as anisotropic super-resolution (where different upsampling factors are needed for each axis) and multi-task learning (e.g., simultaneously upsampling an image and its segmentation mask).

Function-specific documentation is here.

New in Keras 3 Integration

• Dual Backend Support: Seamlessly switch between PyTorch and TensorFlow.
• Universal Anisotropic / Asymmetric Upsampling: All 2D and 3D architectures (ESPCN, LDBPN, WDSR, RCAN, CARN, SRFBN, SAN, AS-DBPN) natively support anisotropic scaling factors (e.g. factor=(1, 2) for 2D or factor=(1, 1, 2), factor=(2, 2, 4) for 3D), ideal for converting thick-slice clinical MRI acquisitions into isotropic high-resolution volumes.
• 3D ESPCN Architecture: High-performance super-resolution using Pixel Shuffling, optimized for Apple Silicon (MPS).
• Patch-wise Inference: Memory-efficient inference with Gaussian blending to eliminate stitching artifacts.
• Blind Perceptual Training: Synthetic data simulation for robust, general-purpose MRI enhancement.
• Low-Latency AS-DBPN: A highly optimized dual-recurrent back-projection model designed for fast, high-fidelity blind super-resolution.
• Live Convergence Reporting & Synchronized 4-Way Comparison: Dark-mode responsive visual reporting dashboard (asdbpn_3d_report.html) featuring an identical-scale 4-way visual comparison grid (Original HR, Downsampled LR at matching FOV, Linear Upsampled, and 3D AS-DBPN SR), unified stacked orthogonal cross-sections, difference residual heatmaps, and strictly monotonic SVG convergence tracking across multi-stage training curriculums (Warmup, Adaptation, Robustness, Joint Fine-Tuning). Supported by a lightweight guardian watcher (tests/watch_report.py) ensuring zero concurrency corruption and sub-20ms instant report compilation.

Key Innovations of the AS-DBPN Network

The Attention-Guided Shared Deep Back-Projection Network (AS-DBPN) is a state-of-the-art super-resolution architecture that achieves top-tier quality metrics while maintaining a highly optimized runtime profile:

1. Lightweight Recurrent Back-Projection: Uses a shared-weight recurrent loop with $T=4$ steps to iteratively project features between low- and high-resolution spaces, providing a parameter-efficient alternative to deep feedforward networks.
2. Simplified Hybrid Attention: Replaces computationally heavy inner-loop attention blocks with a single Socrat/Channel Attention (SOCA/RCAN) block at the final reconstruction layer. This eliminates redundant operations, resulting in a 4.4x speedup / 77% latency reduction (MPS latency reduced from 261.13 ms to 59.05 ms).
3. Shared Layer Normalization: Introduces shared LayerNormalization inside the recurrent loop to stabilize scaling dynamics and prevent gradient collapse during iterative feedback steps.
4. Generalization Under Blind Conditions: Trained via a curriculum of mixed geometries, Rician noise, and out-of-focus blur, the model is resilient to blind, ill-posed degradations, achieving 23.70 dB PSNR on the r16 brain validation patch (+0.84 dB over SAN) and ranking #3 overall on the Mixed Simulation Class benchmark.

Evaluations & Visual Reports

GitHub does not render standalone .html files inline, so the interactive reports below are linked through htmlpreview.github.io, which renders them straight from this repo:

• 9-Class Simulation Benchmark: Overall Model Comparison — interactive PSNR/SSIM/GMSD/HFEN/Correlation comparison across all 12 models (11 architectures + bilinear baseline) and 9 synthetic simulation classes (brain, vessels, fractal noise, sinewave, etc.).
• 3D 9-Class Simulation Benchmark Report — interactive multi-model 3D comparison report across architectures and 3D synthetic classes.
• 3D AS-DBPN Visual Convergence & Quality Dashboard — live auto-updating dashboard monitoring 3D AS-DBPN convergence with identical-scale 4-way comparison (Original, Downsampled, Linear, SR), orthogonal slice montages (Axial, Coronal, Sagittal), residual error maps ($|\text{SR} - \text{GT}|$), and SVG sparklines.
• r16 Brain MRI Qualitative Comparison — side-by-side visual comparison of ground truth, bilinear, SAN, and AS-DBPN reconstructions on the classic r16 test image.
• Simulated Example Gallery — low-res/high-res pairs across all 9 simulation classes used for blind training and evaluation.

Static summaries of the same benchmark are also available as Markdown tables: Class Performance Summary and Rank-of-Ranks & SRFBN Analysis.

What the current results show

• Rank of Ranks (9-class benchmark, full retrain): REF-DBPN ranks #1 overall (score 2.98), followed by SAN (#2, 3.67) and LDBPN (#3, 4.22). All 11 trained architectures beat the Bilinear baseline (10.58) by a wide margin.
• AS-DBPN (the new low-latency model) trades a small amount of rank-score for a 4.4x speedup: MPS latency drops from 261.13 ms to 59.05 ms, while still beating SAN by +0.84 dB PSNR on the r16 validation patch and landing #3 overall on the mixed-class benchmark — the best quality-per-millisecond tradeoff in the collection.
• Heavier back-projection models (REF-DBPN, LDBPN) consistently top the accuracy leaderboard but cost far more latency (REF-DBPN: 648 ms vs. ~15–23 ms for the lightweight models), which motivated the AS-DBPN design.
• SRFBN bilinear-bypass fix: SRFBN previously collapsed to an MSE-minimizing bilinear shortcut. Zero-initializing its LearnableScale skip parameter forced the recurrent conv branch to train from a genuine 15.53 dB start; the retrained model now shows fully active convolutional weights and 8–13% HFEN improvement (sharper high-frequency detail) on structured classes like sinewave, layered, and fractal_noise.

Key Features

SIQ is more than just an inference tool; it's a complete framework that facilitates:

• Flexible Data Generation: Automatically create paired low-resolution and high-resolution patches for training with siq.image_generator. Includes support for multi-channel data (e.g., image + segmentation) via siq.seg_generator.
• Advanced Model Architectures: Easily instantiate powerful Deep Back-Projection Networks (DBPN) for 2D and 3D with siq.default_dbpn, customized for any upsampling factor and multi-task outputs.
• Perceptual Loss Training: Go beyond simple Mean Squared Error. SIQ includes tools for using pre-trained feature extractors (siq.get_grader_feature_network, siq.pseudo_3d_vgg_features_unbiased) to optimize for perceptual quality.
• Intelligent Loss Weighting: Automatically balance complex, multi-component loss functions (e.g., MSE + Perceptual + Dice) with a single command (siq.auto_weight_loss, siq.auto_weight_loss_seg) to ensure stable training.
• End-to-End Training Pipelines: Train models from start to finish with the high-level siq.train and siq.train_seg functions, which handle data generation, validation, and model saving.
• Robust Inference: Apply your trained models to new images with siq.inference, including specialized logic for region-wise and blended super-resolution when guided by a segmentation mask.
• Comprehensive Evaluation: Systematically benchmark and compare model performance with siq.compare_models, which calculates PSNR, SSIM, and Dice metrics against baseline methods.

Installation

You can install the official release directly from PyPI:

pip install siq

To install the latest development version from this repository:

git clone https://github.com/stnava/siq.git
cd siq
pip install .

Quick Start: A 5-Minute Example

The examples demonstrate the core workflow: training a model on publicly available data and using it for inference.

tests/test.py
tests/test_seg.py

Pre-trained Models and Compatibility

We provide a collection of pre-trained models to get you started without requiring you to train from scratch.

• Download Pre-trained Models from Figshare

Important Note on Keras/TensorFlow Versions

The deep learning ecosystem evolves quickly. Models saved with older versions of TensorFlow/Keras (as .h5 files) may have trouble loading in newer versions (TF >= 2.16) due to the transition to the .keras format.

If you encounter issues loading a legacy .h5 model, we provide a robust conversion script. This utility will convert your old .h5 files into the modern .keras format.

Usage:

import siq

# Define the directory containing your old .h5 models
source_dir = "~/.antspymm/" # Or wherever you downloaded the models
output_dir = "./converted_keras_models"

# Convert the models
siq.convert_h5_to_keras_format(
    search_directory=source_dir,
    output_directory=output_dir,
    exclude_patterns=["*weights.h5"] # Skips files that are just weights
)

After running this, you can load the converted models from the converted_keras_models directory using siq.read_srmodel or tf.keras.models.load_model.

For Developers

Setting Up the Environment

This package is tested with Python 3.11 and TensorFlow 2.17. For optimal CPU performance, especially on Linux, you may want to set these environment variables:

export TF_ENABLE_ONEDNN_OPTS=1
export ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS=8
export TF_NUM_INTRAOP_THREADS=8
export TF_NUM_INTEROP_THREADS=8

Note: tests/train_model_refinement.py and generate_summary_images.py force KERAS_BACKEND=torch at import time (for MPS/CUDA acceleration via PyTorch), so refining/training a model requires torch to be installed in addition to the requirements.txt dependencies:

pip install torch

Model Refinement & Fine-Tuning

To refine and fine-tune pre-trained models using our advanced mixed-modality simulation engine (supporting brain structures, sinewaves, layered strips, Rician noise, and coordinate zoom):

# Refine the Channel Attention ESPCN model (default: batch size 1)
python tests/train_model_refinement.py espcn --batch-size 1

# Refine the Lightweight DBPN model
python tests/train_model_refinement.py ldbpn --batch-size 2

# Refine the Reference DBPN model
python tests/train_model_refinement.py ref-dbpn --batch-size 1

Every model accepts --dim {2,3} (default: 3), so 3D training needs no extra flags:

# Refine the 3D AS-DBPN model (low-latency blind super-resolution)
python tests/train_model_refinement.py asdbpn --dim 3 --batch-size 1

# Equivalent to the above, since --dim defaults to 3
python tests/train_model_refinement.py asdbpn --batch-size 1

espcn-rc and wdsr-rc are 2D-only (resize-conv checkerboard-mitigation pilots) and will raise an error if run with --dim 3; every other model type (espcn, ldbpn, ref-dbpn, wdsr, rcan, carn, srfbn, san, asdbpn) supports both dimensionalities.

The script executes a curriculum-based training sequence across 3 stages:

1. Stage 1: Adaptation Phase (Clean Mixed Geometries): 100 warmup iterations on clean mixed geometries.
2. Stage 2: Robustness Fine-Tuning Phase (Low LR + Noise): 150 joint fine-tuning iterations introducing Rician noise.
3. Stage 3: Dedicated Refinement Phase (High-Fidelity Brain Focus): 200 iterations focusing on high-fidelity brain anatomy structures with a very low learning rate ($1 \times 10^{-7}$).

Live Convergence Dashboard & Checkpointing: When running refinement (e.g. for asdbpn), the training pipeline dynamically saves regular convergence checkpoints and auto-renders a visual HTML dashboard (asdbpn_3d_report.html) complete with orthogonal slice montages (Axial, Coronal, Sagittal), residual error maps ($|\text{SR} - \text{GT}|$), and SVG sparklines:

# Refine 3D AS-DBPN with live visual reporting every 25 iterations
python tests/train_model_refinement.py asdbpn --dim 3 \
  --checkpoint-freq 25 \
  --eval-freq 25 \
  --projection-kernel-size 6 \
  --warmup-max-iter 150

• --checkpoint-freq <int>: Frequency in iterations to write asdbpn_3d_step_{iter:04d}.keras and slice montages (default: 50).
• --eval-freq <int>: Frequency in iterations to evaluate validation PSNR/SSIM during warmup (default: 20).
• --projection-kernel-size <int>: Deconv/conv back-projection kernel size (default: 6 for AS-DBPN, matching 2D benchmark capacity).
• Fast procedural simulations: 3D tubular vessel generation is accelerated 118x via scipy.spatial.cKDTree bounded Euclidean queries.

Intelligent Restart & Skipping:

• If a refined model checkpoint already exists (e.g., espcn_3d_attention_refined.keras, ldbpn_3d_refined.keras, or ref_dbpn_3d_refined.keras), the script automatically loads it, skips Stage 1 and Stage 2, and proceeds directly to Stage 3 (Dedicated Refinement).
• To train from scratch starting from the baseline model, delete or rename the existing refined .keras file in the workspace.

Generating Evaluation Reports

generate_summary_images.py builds the interactive HTML comparison report (see Evaluations & Visual Reports) from whichever refined checkpoints are present in the repo root. It also accepts --dim:

# 2D report (default): reads *_2d_refined.keras, writes summary_results.html
python generate_summary_images.py

# 3D report: reads *_3d_refined.keras, writes summary_results_3d.html
python generate_summary_images.py --dim 3

Any model without a matching checkpoint for that dimensionality shows up as "TBD" in the report rather than causing an error, so it's safe to run before every model has been trained in 3D.

Publishing a New Release

To publish a new version of siq to PyPI:

# Ensure build and twine are installed
python -m pip install build twine

# Clean previous builds
rm -rf build/ siq.egg-info/ dist/

# Build the package
python -m build .

# Upload to PyPI
python -m twine upload --repository siq dist/*