Add "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models" - #192
Open
BrianPengT wants to merge 19 commits into
Open
Add "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models"#192BrianPengT wants to merge 19 commits into
BrianPengT wants to merge 19 commits into
Conversation
layer.py: can choose to use einsum or bmm for self-attention. loss.py: can pass a generator during forward for reproducibility. sampling.py: use vector_norm instead of norm for multiple dimensions. utils.py: added some utility functions for training and inference. trainer.py: deprecated.
… adjoint for hijacking.
There was a problem hiding this comment.
Pull request overview
This PR adds a new “score_inverse” model implementation and accompanying train/validation/test scripts for score-based inverse problems in CT (targeting LIDC-IDRI), alongside supporting updates to CT utilities, classical reconstruction baselines, and image-quality metrics.
Changes:
- Introduces
LION/models/score_inverse/(SDEs, NCSN++ model, sampling/hijack utilities, EMA, losses, and operators) plus runnable scripts underscripts/score_inverse_scripts/. - Updates LIDC-IDRI loader to support an
"image_only"task and improves slice selection / CPU fallback behavior. - Fixes/extends several utilities and metrics (noise utility robustness, PSNR/SSIM handling, HAARPsi init, classical algorithms batching) and adds an FBP baseline.
Reviewed changes
Copilot reviewed 24 out of 25 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/score_inverse_scripts/train.py | New distributed training script for NCSN++ score matching on LIDC-IDRI. |
| scripts/score_inverse_scripts/validation.py | New validation runner to evaluate checkpoint loss deterministically. |
| scripts/score_inverse_scripts/test.py | New reconstruction/evaluation runner with PC sampling + data-consistency hijacking. |
| scripts/score_inverse_scripts/configs.py | New geometry/noise/SDE config presets and paths for scripts. |
| LION/models/score_inverse/utils.py | New seed/checkpoint/batching utilities for score-inverse stack. |
| LION/models/score_inverse/sde.py | New forward/reverse SDE abstractions + VESDE implementation. |
| LION/models/score_inverse/loss.py | New score-matching loss (SMLoss) with generator support. |
| LION/models/score_inverse/ncsnpp.py | New NCSN++ model implementation and score_fn wrapper. |
| LION/models/score_inverse/layer.py | New building blocks (FIR up/downsample, attention, resblocks, etc.). |
| LION/models/score_inverse/sampling.py | New PC samplers + hijack functions for data consistency and conditional scoring. |
| LION/models/score_inverse/ema.py | New EMA helper for training/eval. |
| LION/models/score_inverse/fst.py | New Fourier Slice Theorem radon implementation for parallel-beam. |
| LION/models/score_inverse/sirt_adj.py | New SIRT-style preconditioned adjoint pseudo-inverse operator. |
| LION/models/score_inverse/init.py | Package init for new score_inverse module. |
| LION/models/score_inverse/.gitignore | Local ignore entry for score_inverse directory. |
| LION/CTtools/ct_utils.py | Improves sinogram noise function device/shape robustness. |
| LION/data_loaders/LIDC_IDRI.py | Adds image_only task, CPU fallback, and “all slices” loading behavior. |
| LION/metrics/ssim.py | Makes SSIM more robust to squeezing/batching + adds optional data_range. |
| LION/metrics/psnr.py | Adds optional data_range and adjusts batching/squeezing logic. |
| LION/metrics/haarpsi.py | Adds missing super().__init__() and updates “Last update” date. |
| LION/classical_algorithms/tv_min.py | Fixes batched input handling (per-batch call). |
| LION/classical_algorithms/sirt.py | Fixes batched input handling (per-batch call). |
| LION/classical_algorithms/fbp.py | Adds an FBP baseline wrapper with batched support. |
| LION/classical_algorithms/init.py | Exposes fbp in the classical algorithms API. |
| CHANGES.md | New changelog-style summary describing the fork’s changes. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
26
to
28
| def forward( | ||
| self, x: torch.Tensor, target: torch.Tensor, reduce=str | None, batched=True | ||
| self, x: torch.Tensor, target: torch.Tensor, reduce=str | None, batched=True, data_range: float | None = None | ||
| ) -> torch.Tensor: |
Comment on lines
+11
to
+18
| def set_global_seed(seed=0): | ||
| random.seed(seed) | ||
| np.random.seed(seed) | ||
| torch.manual_seed(seed) | ||
| torch.cuda.manual_seed_all(seed) | ||
|
|
||
| torch.backends.cudnn.benchmark = False | ||
| torch.use_deterministic_algorithms(True, warn_only=True) |
Comment on lines
+159
to
+168
| torch.cuda.set_device(local_rank) | ||
| set_global_seed(args.seed) | ||
|
|
||
| train_loader, train_sampler = get_dataloader( | ||
| CONFIG, args.data_prop, args.seed, local_rank, world_size, is_distributed | ||
| ) | ||
|
|
||
| sde = VESDE(sigma_min=CONFIG['sigma_min'], sigma_max=CONFIG['sigma_max']) | ||
|
|
||
| model = NCSNpp(**CONFIG['model_kwargs']).to(local_rank) |
| for batch in train_loader: | ||
| optimizer.zero_grad(set_to_none=True) | ||
|
|
||
| x = from_HU_to_normal(batch.to(local_rank, non_blocking=True)) |
Comment on lines
+37
to
+53
| parser = argparse.ArgumentParser() | ||
| parser.add_argument('--checkpoint_path', type=str, nargs='+', default=['/home/tp534/rds/hpc-work/trained_models/ncsnpp/checkpoint_epoch_199_step_2376399.pth'], help='Path to the checkpoint file(s)') | ||
| parser.add_argument('--device', type=str, default='cuda', help='Device to run validation on (GPU is assumed and required)') | ||
| parser.add_argument('--lb', type=float, nargs='+', default=[0.841], help='Parameter lambda for hijacking') | ||
| parser.add_argument('--N', type=int, nargs='+', default=[100], help='Number of predictor steps') | ||
| parser.add_argument('--M', type=int, default=1, help='Number of corrector steps') | ||
| parser.add_argument('--snr', type=float, default=0.246, help='Corrector SNR') | ||
| parser.add_argument('--use_old_sampler', action='store_true', help='If set, use the old pc_sampler instead of pc_sampler_new') | ||
| parser.add_argument('--output_dir', type=str, default=data_dir, help='Directory to save the output reconstructions') | ||
| parser.add_argument('--output_name', type=str, default='score_inv', help='Reconstruction name for output saving') | ||
| parser.add_argument('--compile', action='store_true', help='If set, compile the model using torch.compile') | ||
| parser.add_argument('--bf16', action='store_true', help='Use bfloat16 precision for validation (only effective on CUDA devices)') | ||
| parser.add_argument('--sinogram_path', type=str, default=None, help='Path to precalculated .npy sinogram file') | ||
| parser.add_argument('--pseudo_inv_mode', type=str, default='sirt', choices=['fst', 'sirt'], help='Reconstruction operator to use for pseudo-inverse') | ||
| parser.add_argument('--geometry', type=str, default='parallel', choices=['parallel', 'fan'], help='CT geometry type') | ||
| parser.add_argument('--clean_hijack', action='store_true', help='If set, use clean hijacking (no noise injection to the sinogram in hijack function)') | ||
| parser.add_argument('--seed', type=int, default=0, help='Random seed for reproducibility') |
|
|
||
| # Seed global RNGs and enforce deterministic algorithms so dataset and | ||
| # sampling are reproducible across runs. | ||
| set_global_seed(0) |
Comment on lines
+14
to
+20
| if sino.dim() == 4: | ||
| B, _, _, _ = sino.shape | ||
| remove_batch = False | ||
| elif sino.dim() == 3: | ||
| B = 1 | ||
| sino = sino.unsqueeze(0) | ||
| remove_batch = True |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Details see CHANGES.md
Added "Solving Inverse Problems in Medical Imaging with Score-Based Generative Models" to the model and added training/validation/testing scripts using the LIDC-IDRI dataset.
Did minor debugging and updates to CT utility functions and metrics.