MaLViL: Multi-axis Low-rank Vision-LSTM for Medical Image Segmentation

Official pretrained weights for MaLViL (MLMI Workshop, MICCAI 2026), a 2D segmentation network that keeps Vision-LSTM in every decoder stage, for abdominal CT, breast ultrasound and dermoscopy images.

Quick start

You do not need to clone anything. PyTorch Hub fetches the model code from GitHub, and the weights are downloaded from this repository.

pip install torch torchvision "timm>=1.0.15" einops "numpy<2" scipy SimpleITK huggingface_hub

Segment your own images

The predictor applies the training preprocessing and returns a mask with the size of your input.

import torch

# Skin lesion (dermoscopy image)
skin = torch.hub.load("xmindflow/MaLViL", "predictor", dataset="isic2018", trust_repo=True)
mask = skin.predict("lesion.jpg")            # (H, W): 0 = background, 1 = lesion

# Abdominal CT (NIfTI volume, processed slice by slice)
ct = torch.hub.load("xmindflow/MaLViL", "predictor", dataset="synapse", trust_repo=True)
mask = ct.predict("abdomen_ct.nii.gz")       # (D, H, W)
print(ct.classes)                            # label names, index = label value

predict accepts a file path (.png, .jpg, .npy, .nii, .nii.gz), a PIL image or a NumPy array. It runs on the GPU if one is available.

Use the model directly

model = torch.hub.load("xmindflow/MaLViL", "malvil", dataset="synapse", trust_repo=True)
x = torch.randn(1, 1, 224, 224)              # one preprocessed slice
with torch.no_grad():
    (logits,), rec_error = model(x)          # logits: (1, 9, 224, 224)
pred = logits.argmax(dim=1)

Train on your own data

malvil_custom returns an untrained network with an ImageNet PVTv2-B2 encoder, for your own training loop. The network is built for one input size, a multiple of 32.

net = torch.hub.load("xmindflow/MaLViL", "malvil_custom", in_chs=3, n_classes=4, img_size=224, trust_repo=True)
(logits,), rec_error = net(images)
loss = criterion(logits, labels) + 0.01 * rec_error

The repository also has a ready training and evaluation script that takes a folder of images and masks.

Models

dataset	Task	Input	Labels	File
synapse	Abdominal CT, 8 organs	1 × 224 × 224	0 background, 1 aorta, 2 gallbladder, 3 left kidney, 4 right kidney, 5 liver, 6 pancreas, 7 spleen, 8 stomach	malvil-synapse.pth
busi	Breast ultrasound	1 × 256 × 256	0 background, 1 lesion	malvil-busi.pth
ham10000	Skin lesion	3 × 224 × 224	0 background, 1 lesion	malvil-ham10000.pth
ph2	Skin lesion	3 × 224 × 224	0 background, 1 lesion	malvil-ph2.pth
isic2017	Skin lesion	3 × 224 × 224	0 background, 1 lesion	malvil-isic2017.pth
isic2018	Skin lesion	3 × 224 × 224	0 background, 1 lesion	malvil-isic2018.pth

Each file is a plain PyTorch state_dict (about 150 MB). The network has size-dependent parameters, so BUSI weights only load at 256 and the others only at 224.

Preprocessing

The predictor does this for you. If you call the model directly, prepare inputs the same way:

• Synapse (CT): axial slices. Clip Hounsfield units to [-125, 275], scale to [0, 1], resize to 224 × 224 (cubic).
• BUSI (ultrasound): grayscale image. Resize to 256 × 256 (bilinear), then min-max normalize the image to [0, 1].
• HAM10000 / PH2 / ISIC 2017 / ISIC 2018 (dermoscopy): RGB image. Resize to 224 × 224 (bilinear), then min-max normalize the image to [0, 1].

Predictions are resized back to the input size with nearest-neighbor interpolation.

Results

These are retrained checkpoints, so they differ slightly from the tables in the paper. Each row was measured with scripts/main.sh DATASET TEST from the repository. All use AdamW and a learning rate of 1e-4.

Dataset	Input	Epochs	Batch size	Seed	Loss	Dice	Acc	IoU
PH²	224	150	8	2222	dice+ce	95.78	97.54	92.06
HAM10000	224	20	16	1234	dice+ce	95.15	97.51	91.40
ISIC 2017	224	40	8	1234	dice+ce	92.13	97.14	86.22
ISIC 2018	224	40	8	1234	dice+ce	91.08	96.53	84.92
BUSI	256	50	8	1234	dice+ce	82.38	96.92	74.33
Synapse	224	350	8	42	boundary	85.68 (HD95 17.01)

Download the weights only

hf download afshinbigboy/MaLViL --include "*.pth" --local-dir weights

or in Python:

from huggingface_hub import hf_hub_download
path = hf_hub_download(repo_id="afshinbigboy/MaLViL", filename="malvil-synapse.pth")

Intended use and limitations

These models are for research only. They are not a medical device and must not be used for diagnosis or treatment decisions. Each model was trained on a single public dataset and may perform poorly on images from other scanners, protocols or populations.

License

The code and weights are released under the MIT license. The training datasets have their own terms (for example, HAM10000 is CC BY-NC 4.0), so please also respect the license of each dataset.

Citation

If you use MaLViL, please cite:

@inproceedings{bozorgpour2026malvil,
  title         = {{MaLViL}: Multi-axis Low-rank Vision-{LSTM} for Medical Image Segmentation},
  author        = {Bozorgpour, Afshin and Ghorbani Kolahi, Sina and Heidari, Moein and Hacihaliloglu, Ilker and Merhof, Dorit},
  booktitle     = {Medical Image Computing and Computer Assisted Intervention -- MICCAI 2026 Workshops and Challenges (MLMI)},
  series        = {Lecture Notes in Computer Science},
  volume        = {17272},
  publisher     = {Springer Nature Switzerland},
  year          = {2026},
  eprint        = {2608.17635},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV}
}