You can deploy this model and start earning money today!

DRDMsig/
Omini3D

Name: Omini3D
Rating: 2.5 (881 reviews)

Views

No views yet

OmniMorph

Deform All-in-One Framework for Medical Image Generation, Restoration and Registration based on a conditional Deformation-Recovery Diffusion Model (DeformDDPM).

OmniMorph is a unified framework for 2D/3D multi-modal medical imaging (CT, MRI, PET) supporting:

Generation — text-conditioned image synthesis via BERT embeddings.
Restoration — recover anatomically plausible images from degraded inputs.
Registration — paired / unpaired / flexible-resolution registration via diffused deformation vector fields.

Repository Contents

Path	Description
`OM_train*.py`	Training entrypoints (single-/2-/3-mode variants, CUDA + Intel XPU)
`OM_aug.py`, `OM_reg.py`, `OM_contrastive*.py`	Inference / augmentation / registration / contrastive scripts
`Diffusion/`	DeformDDPM core: `diffuser.py`, networks, losses, spatial utils
`OMorpher/`	Higher-level model wrapper
`Dataloader/`	Multi-modality dataloaders + dataset mappings (16 datasets)
`Config/`	YAML training/inference configs
`Scripts/`	Auxiliary scripts (registration, evaluation)
`tests/`	Pytest suite for `OMorpher` and loss functions
`bash_.sh`, `.slurm`	SLURM submission scripts (CUDA + Intel XPU/Dawn)
`Models/all_om_net/000110_all_om_net.pth`	Trained checkpoint — production multi-modal `recmulmodmutattnnet` (epoch 110, ~3.0 GB)
`Models/all_recmulmodmutattnnet/000010_all_recmulmodmutattnnet.pth`	Earlier `recmulmodmutattnnet` run (epoch 10, ~906 MB)

Note Only the final checkpoint of each training run is shipped — intermediate epochs and the bert_large_uncased weights are not bundled. Download bert-large-uncased from the official Hugging Face repo if you need the contrastive text encoder.

Setup

bash

git clone https://huggingface.co/DRDMsig/Omini3D
cd Omini3D
pip install -r requirements.txt

For Intel XPU / Dawn cluster, install the matching intel-extension-for-pytorch build before installing the rest of the requirements.

Quick Start

Training

bash

# Single-mode diffusion
CUDA_VISIBLE_DEVICES=0 python OM_train.py -C Config/config_om.yaml

# Dual mode (diffusion + registration)
CUDA_VISIBLE_DEVICES=0,1 python OM_train_2modes.py -C Config/config_om.yaml

# Triple mode (diffusion + contrastive + registration)
CUDA_VISIBLE_DEVICES=0,1 python OM_train_3modes.py -C Config/config_om.yaml

# Intel XPU (single node)
sbatch bash_train_single_node.sh

Inference

bash

# Augmentation / restoration with a trained model
python OM_aug.py -C Config/config_om.yaml

# Paired registration
python OM_reg.py -C Config/config_om.yaml

# Flexible-resolution registration
python OM_reg_flexres.py -C Config/config_om.yaml

Loading the checkpoint

python

import torch
from Diffusion.networks import get_net

# Production network (multi-modal recmulmodmutattnnet)
net = get_net("recmulmodmutattnnet")

# Production checkpoint (epoch 110)
ckpt_path = "Models/all_om_net/000110_all_om_net.pth"
# Or earlier run: "Models/all_recmulmodmutattnnet/000010_all_recmulmodmutattnnet.pth"

state = torch.load(ckpt_path, map_location="cpu")
net.load_state_dict(state["model"] if "model" in state else state)
net.eval()

Architecture

Config YAML → DataLoader(s) → DeformDDPM(Network, STN) → Loss → Checkpoint

DeformDDPM (Diffusion/diffuser.py) — forward/reverse diffusion over deformation vector fields (DVFs); multi-scale DDFs at control-point ratios [4, 8, 16, 32, 64].
Networks (Diffusion/networks.py) — selectable via get_net(name):
- recmulmodmutattnnet — current production multi-modal multi-head-attention net (used by 000110_all_om_net.pth)
- recmutattnnet, recmutattnnet_contrastive, recresacnet, defrecmutattnnet
STN — Spatial Transformer for differentiable warping; composes deformations as comp_ddf = dvf + stn(ddf, dvf).
Losses (Diffusion/losses.py, losses_ncc0.py) — Grad, LNCC, LMSE, NCC, MRSE, RMSE.

Datasets Supported

Dataloader/nifty_mappings/ contains pre-computed mappings for 16 public medical-imaging datasets, including: AbdomenAtlas, AbdomenCT-1k, BraTS 2019/2020/2021, MSD, OASIS-1/2, OAI-ZIB, MnMs, Kaggle OSIC, TotalSegmentator (CT+MRI), PSMA-FDG-PET-CT-Lesion, CIA.

The dataset files themselves are not included; obtain them from their respective sources and update the mapping paths.

Citation

bibtex

@article{omnimorph,
  title  = {OmniMorph: Deform All-in-One Framework for Medical Image Generation,
            Restoration and Registration via Conditional Deformation-Recovery
            Diffusion Models},
  author = {Zheng, J. and Mo, M. and others},
  year   = {2025}
}

License

MIT — see LICENSE.