Home

Awesome

MedIM: One-Line Code for Pre-trained Medical Image Models in PyTorch

x x

A collection of pre-trained medical image models in PyTorch. This repository aims to provide a unified and easy-to-use interface for comparing and deploying these models.

Supported Models

Quick Start

Setup Environment

You can use this cmd to install this toolkit via pip:

pip install medim

For developers, you can install in the editable mode via:

git clone https://github.com/uni-medical/MedIM.git
cd MedIM
pip install -e .

Example Usage

First, let us import medim.

import medim

You have four ways to create a PyTorch-compatible model with create_model:

1. use models without pretraining

model = medim.create_model("STU-Net-S") 

2. use local checkpoint

model = medim.create_model(
            "STU-Net-S",
            pretrained=True,
            checkpoint_path="../tests/data/small_ep4k.model") 

3. use checkpoint pre-trained on validated datasets (will automatically download it from HuggingFace)

model = medim.create_model("STU-Net-B", dataset="BraTS21")

4. use HuggingFace url (will automatically download it from HuggingFace)

model = medim.create_model(
            "STU-Net-S",
            pretrained=True,
            checkpoint_path="https://huggingface.co/ziyanhuang/STU-Net/blob/main/small_ep4k.model") 

Tips: you can use MEDIM_CKPT_DIR environment variable to set custom path for medim model downloading from huggingface.

Then, you can use it as you like.

input_tensor = torch.randn(1, 1, 128, 128, 128)
output_tensor = model(input_tensor)
print("Output tensor shape:", output_tensor.shape)

More examples are in examples.

Roadmap & TO-DO