Skip to content

Prediction Python API

Use dmri predict for the complete folder-based prediction workflow. For Python code that manages model inputs and inference directly, the package root exposes dmri.load_pretrained, dmri.list_pretrained_models, and dmri.PretrainedModel.

import dmri

available = dmri.list_pretrained_models()
bundle = dmri.load_pretrained("msb3s_2_4_6_128")

model = bundle.model
simulator = bundle.simulator
config = bundle.config
step = bundle.step

By default, these functions use the public manugloeck/dmri-pretrained repository. Model files are downloaded and cached by huggingface_hub. load_pretrained accepts revision, cache_dir, local_files_only, and a network-forward precision. CPU fp16 requests fall back to fp32.

load_pretrained applies EMA parameters when available and otherwise applies raw parameters. The returned model is in evaluation mode. The simulator field is the reconstructed MultiCompartment class.

Pretrained models are acquisition-specific. Verify compatibility with the input acquisition before using their output. DMRI has not been clinically validated.

API

hub

PretrainedModel dataclass 

PretrainedModel(model: Any, simulator: Any, config: Any, step: int | None)

model instance-attribute 

model: Any

simulator instance-attribute 

simulator: Any

config instance-attribute 

config: Any

step instance-attribute 

step: int | None

load_pretrained 

load_pretrained(model_name=DEFAULT_MODEL_NAME, repo_id=DEFAULT_REPO_ID, which='best', *, revision=None, cache_dir=None, token=None, local_files_only=False, precision=None) -> PretrainedModel

Load a pretrained DMRI model from the Hugging Face Hub.

Parameters

model_name: Checkpoint name inside the repository. repo_id: Hugging Face model repository identifier. which: Which checkpoint to load ("latest", "best", or a step number). revision: Branch, tag, or commit hash to pin. cache_dir: Local directory for downloaded artifacts. token: Hugging Face access token for private repos. local_files_only: Skip any remote download and use the cache exclusively. precision: Cast parameters to this precision (e.g. "fp16", "bf16").

Returns

PretrainedModel The instantiated model, its simulator, config, and training step.

list_pretrained_models 

list_pretrained_models(repo_id=DEFAULT_REPO_ID, *, revision=None, token=None) -> list[str]

Return the names of pretrained models available in a repository.

Parameters

repo_id: Hugging Face model repository identifier. revision: Branch, tag, or commit hash to inspect. token: Hugging Face access token for private repos.

Returns

list[str] Sorted model names found in the repository.