feat: add adaptive sampling, motion registration, MEDL, and e2e ADS configs - #316
Open
georgeyiasemis wants to merge 127 commits into
Open
feat: add adaptive sampling, motion registration, MEDL, and e2e ADS configs#316georgeyiasemis wants to merge 127 commits into
georgeyiasemis wants to merge 127 commits into
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
georgeyiasemis
changed the title
…onfig LPD/VarSplit/ViT engines were refining coil sensitivities twice via _do_iteration and forward_function. ConjGrad Calgary configs need unet_conv_out_bias so Hub weights with conv_out.bias load correctly. Co-authored-by: Cursor <cursoragent@cursor.com>
Zoo validation metrics used crop_outer_slices (slice 50:-50). Released Hub YAMLs omitted it, so full-volume eval scored outer slices. Mirror configs under projects/calgary_campinas/configs/inference and keep the flag in minimize_inference_yaml. Co-authored-by: Cursor <cursoragent@cursor.com>
Add an inference block with crop_outer_slices and commented 5x/10x accelerations to each base_*.yaml instead of separate Hub-style YAMLs. Co-authored-by: Cursor <cursoragent@cursor.com>
Fill missing Unet/RIM/VarNet/LPD/RecurrentVarNet fields from the trained configs and add 4x inference with commented 8x masking for Hub-style predict. Co-authored-by: Cursor <cursoragent@cursor.com>
Rename AXT1 ablation to t6_h128_nl2, drop fastmri metrics, set inference batch_size to 10, and sync Calgary VarNet comparison params. Co-authored-by: Cursor <cursoragent@cursor.com>
6 tasks
georgeyiasemis
requested review from
jonasteuwen
and
a lite review from Copilot
August 11, 2026 13:09
jonasteuwen
reviewed
Aug 11, 2026
jonasteuwen
left a comment
jonasteuwen
left a comment
Contributor
There was a problem hiding this comment.
Please look at my comments
Comment on lines
+40
to
+43
|
|
||
| for i in range(batch_size): | ||
| x = inputs[i : i + 1] | ||
| count = 0 |
Contributor
There was a problem hiding this comment.
Can you do @jit.compile? Or doesn't that work here
Contributor
Author
There was a problem hiding this comment.
No because of the rejection sampling
Comment on lines
+91
to
+96
| super().__init__() | ||
| self.slope = slope | ||
| self.clamp = clamp | ||
|
|
||
| self.fun = ThresholdSigmoidMaskFunction.apply | ||
|
|
Contributor
There was a problem hiding this comment.
I prefer self._variable for class private variables
Comment on lines
-49
to
-50
| data["sensitivity_map"] = self.compute_sensitivity_map(data["sensitivity_map"]) | ||
|
|
Contributor
Author
There was a problem hiding this comment.
This was a bug. Sens estimation is already happening in main engine
Comment on lines
+299
to
+301
| self.image_initialization = ( | ||
| image_initialization.lower() if isinstance(image_initialization, str) else image_initialization | ||
| ) |
Contributor
There was a problem hiding this comment.
Hmm, feels like you should override the __eq__ in the object and make sure str doesn't end up here.
Comment on lines
-67
to
-68
| data["sensitivity_map"] = self.compute_sensitivity_map(data["sensitivity_map"]) | ||
|
|
Contributor
Author
There was a problem hiding this comment.
This was a bug. Sens estimation is already happening in main engine
Comment on lines
+26
to
+38
| def _deform_random_grid(*args, **kwargs) -> list[np.ndarray]: | ||
| """Lazy import wrapper for :func:`elasticdeform.deform_random_grid`. | ||
|
|
||
| ``elasticdeform`` wheels may be incompatible with NumPy 2.x. | ||
|
|
||
| Returns | ||
| ------- | ||
| list[np.ndarray] | ||
| Deformed images returned by ``elasticdeform``. | ||
| """ | ||
| try: | ||
| from elasticdeform import deform_random_grid | ||
| except ImportError as exc: |
Contributor
There was a problem hiding this comment.
Isn't this a function that can be forked/vendored?
Comment on lines
+346
to
+348
| curr_model_cfg.update(operator_kwargs) | ||
| # Drop keys unknown to the constructor (e.g. ModConv defaults on legacy modules). | ||
| additional_models[k] = curr_model(**filter_arguments_by_signature(curr_model, curr_model_cfg)) |
Contributor
There was a problem hiding this comment.
Can this cause silent bugs?
| ARG PYTHON="3.12" | ||
|
|
||
| # TODO: conda installs its own version of cuda | ||
| FROM nvidia/cuda:${CUDA}-devel-ubuntu18.04 |
Contributor
There was a problem hiding this comment.
Probably need to boost this too
Comment on lines
49
to
51
| "pytest>=9", | ||
| "ruff>=0.16.1", | ||
| "ruff>=0.14", | ||
| "ty>=0.0.38", |
Contributor
There was a problem hiding this comment.
I would keep it to 0.16.1 or you run into issues.
Prefer InitType for RIM init, private binarizer attrs, warn on dropped model kwargs, vendor SciPy elastic deform (drop elasticdeform), pin ruff 0.16.1, and boost the Docker CUDA/PyTorch stack. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Remove `from __future__ import annotations` from Python files touched in this PR and delete the obsolete `projects/scripts/minimize_inference_yaml.py` helper from the repository. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Predict without targets when inference.metrics is empty, only write metrics_inference.json when metrics are computed, and fix DirectEnum from_str typing after dropping future annotations. Co-authored-by: Cursor <cursoragent@cursor.com>
Align the annotation with the (volumes, metrics) tuple returned by inference, and silence BaseConfig/subclass attribute checks. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Keep ADS as an example rather than baking acquisition-specific wording into the H5 writer API docs. Co-authored-by: Cursor <cursoragent@cursor.com>
Drop branch-only batch_size bumps and redundant regularizer_dropout defaults that did not change behavior. Co-authored-by: Cursor <cursoragent@cursor.com>
Restore validation metrics and main inference batch sizes, remove redundant model defaults (dropouts, unused RIM initializer fields), and keep only real diffs: inference blocks and SENSE enum casing. Co-authored-by: Cursor <cursoragent@cursor.com>
Inference lives on the training YAMLs (4x active, 8x commented); the rate-specific copies were duplicates. Update the project README accordingly. Co-authored-by: Cursor <cursoragent@cursor.com>
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.
Description
This PR brings end-to-end adaptive dynamic sampling (ADS) and optional motion registration into DIRECT, together with the paper experiment layouts for the MIDL 2026 e2e ADS works.
It wires learned sampling policies and registration models through the MRI training/inference stack, ships validated train + inference YAML pairs under
projects/e2e_ads_reconandprojects/e2e_ads_recon_reg, and hardens config/checkpoint loading so those projects (and Hub-style inference YAMLs) run cleanly on currentmain(including ruff 0.16).What's new
direct/nn/adaptive/): policies for 1D/2D and static/dynamic masks, rejection-sampling as a typed training path (RejectionSamplingError), and budgeting that works with static ACS on dynamic volumes.direct/registration/,direct/nn/registration/): VoxelMorph-style learned registration plus classical utilities (demons, optical flow, elastic deformation); joint or decoupled recon↔registration training.source_key/target_key) and displacement-field visualization for registration runs.projects/e2e_ads_recon— adaptive sampling + reconstruction (vSHARP / MEDL variants, frame/init2 configs).projects/e2e_ads_recon_reg— same family with registration (vSHARP / VarNet / LOUPE / fixed-mask variants, phase/frame and disjoint options).*_inference.yaml): one inference file per model with active val-4× masking and commented 6×/8× alternatives (no per-rate file explosion).What's changed
MRIModelEngine) call into adaptive sampling / registration when configured; writers can persist registration volumes and displacement fields.training/validationare optional so inference-only YAMLs are valid;logging.log_intervalcontrols TensorBoard/scalar flush cadence.RuntimeErroron missing keys; ModConvconv_out_bias/ AdaIN kwargs wiring aligned so paper weights load under current UNet construction.filter_arguments_by_signaturekeeps**kwargs-based model builders working when configs carry extraimage_*fields.mainruff cleanup without dropping the above behavior.User impact
training/validationblocks; use*_inference.yaml(or the minimizer) instead of maintaining many rate-specific files.registration_simulate_referencein transforms); without that, predict fails with a missingreference_imagekey.