UMH

Official repository for UMH: Enhancing Medical Image Segmentation via Heat Conduction Equation.

Installation

Requirements: Ubuntu 22.04, CUDA 12.8

1. Create a virtual environment: conda create -n UMH python=3.11 -y and conda activate UMH
2. Install Pytorch 2.4.1: pip install torch==2.4.1 torchvision==0.19 --index-url https://download.pytorch.org/whl/cu128
3. Install Mamba: pip install causal-conv1d>=1.2.0 and pip install mamba-ssm --no-cache-dir
4. Download code: git clone https://github.com/Rows21/UMH
5. cd UMH/umh and run pip install -e .

sanity test: Enter python command-line interface and run

import torch
import mamba_ssm

Model Training

UMH is built on nnU-Net. please follow this guideline to prepare your own dataset. The dataset used in the paper is here.

Preprocessing

nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity

Train 2D models

• Train 2D U-Mamba_Bot model

nnUNetv2_train DATASET_ID 2d all -tr nnUNetTrainerUMH

Train 3D models

• Train 3D U-Mamba_Bot model

nnUNetv2_train DATASET_ID 3d_fullres all -tr nnUNetTrainerUMH

Inference

• Predict testing cases with UMH model

nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_ID -c CONFIGURATION -f all -tr nnUNetTrainerUMH --disable_tta

``

CONFIGURATION can be 2d and 3d_fullres for 2D and 3D models, respectively.

Paper

@article{wu2025enhancing,
  title={Enhancing Medical Image Segmentation via Heat Conduction Equation},
  author={Wu, Rong and Yu, Yim-Sang},
  journal={arXiv preprint arXiv:2511.03260},
  year={2025}
}

Acknowledgements

We acknowledge all the authors of the employed public datasets, allowing the community to use these valuable resources for research purposes. We also thank the authors of nnU-Net and U-Mamba for making their valuable code publicly available.