Skip to content

Using a pre-trained model (local inference)

MicroStructFormer models are provided as local model packages and can be used directly for inference without accessing any online repositories. This section describes the minimal steps required to load a pre-trained model from disk and run inference on electron microscopy images.

1. Install dependencies

Configure a Python environment and install the required dependencies from the command line.
MicroStructFormer requires Python 3.11 or 3.12.

You may use any environment management tool of your choice (e.g. Conda, venv, or a system Python).
The following example illustrates a setup using Conda: 

conda create --name MSFenv python=3.11 # or python=3.12
conda activate MSFenv

And then install necessary packages:

pip install numpy
pip install torch torchvision
pip install transformers
pip install pillow numpy
pip install scikit-image matplotlib

These dependencies support model inference, mask post-processing, and visualization utilities. If you plan to run inference on a GPU, please ensure that PyTorch is installed with a CUDA build appropriate for your system.
For detailed installation instructions, refer to the official PyTorch documentation.

2. Prepare local paths

[!IMPORTANT] Data and Model Preparation Required Before running the code below, please ensure you have downloaded the required data and model weights from the Model Zoo and Dataset sections. You must modify the MODEL_DIR and IMAGE_DIR variables in the code to point to your local paths before execution.

Before running inference, use the download_model.py script provided in the Model Zoo section to automatically download the pre-trained model package. By default, it will create a folder (e.g., MicroStructFormer_Swin_Large) containing the model weights, configuration files, and processor configuration.

You will also need a local directory containing input images. 

# Local paths
MODEL_DIR = r""   # path to the local model folder you save/create
IMAGE_DIR = r""   # path to images
MODEL_DIR must contain config.json, preprocessor_config.json, and model.safetensors.

3. Load the model and processor

The model and image processor are loaded from the local directory using from_pretrained in offline mode： 

import os
import torch
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_model():
    model = Mask2FormerForUniversalSegmentation.from_pretrained(
        MODEL_DIR, local_files_only=True
    )
    processor = AutoImageProcessor.from_pretrained(
        MODEL_DIR, local_files_only=True
    )
    model.eval().to(device)
    return model, processor

4. Preprocess input images

Input images are loaded from disk and converted to the format expected by the model： 

def preprocess_image(image_path, processor):
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    return inputs.pixel_values.to(device), image

5. Run inference

The following example shows how to run inference on a single image and obtain a semantic segmentation mask. 

model, processor = load_model()

# Example image path
image_name = "" # your image name
image_path = os.path.join(IMAGE_DIR, image_name)

# Load and preprocess image
input_tensor, original_image = preprocess_image(image_path, processor)

# Forward pass
with torch.no_grad():
    output = model(input_tensor)

# Post-process to target size [H, W]
target_size = [original_image.size[1], original_image.size[0]]

# Convert model outputs (query-level predictions) into a pixel-wise
# semantic segmentation map and resize it to the target spatial size.
pred = processor.post_process_semantic_segmentation(
    output, target_sizes=[target_size]
)[0]

pred_np = pred.cpu().numpy().astype(np.uint8)

The resulting pred_np array contains the predicted semantic labels for each pixel and can be used for visualization or downstream quantitative analysis.

6. Visualizing segmentation masks

MicroStructFormer produces semantic segmentation masks in which each pixel stores an integer label ID. These masks can be visualized by mapping label IDs to colors using a simple lookup table.

Because axon–myelin segmentation and subcellular segmentation are produced by separate models, the label space may include compartments originating from different sources.
All label IDs are preserved to give users full flexibility in visualization and downstream analysis.

Label IDs and color mapping

A typical label convention used in this project is shown below:

Label ID Structure
0 background
1 abnormality
2 mitochondria-like organelle
3 mitochondria
4 mitochondria outside axon
5 periaxonal space
6 axon
7 myelin

Depending on your pipeline, some labels (e.g. axon or myelin) may originate from external segmentation tools, while others are predicted by MicroStructFormer.

Predicted mask colorization example code

The following example code shows how to convert a single-channel label mask into an RGB image using a fixed colormap.

import numpy as np
from skimage.io import imread, imsave
from matplotlib.colors import ListedColormap

# Define a colormap (index = label ID)
cmap = ListedColormap([
    "black",    # 0 background
    "red",      # 1 abnormality
    "yellow",   # 2 mito-like organelle
    "green",    # 3 mitochondria
    "cyan",     # 4 mitochondria outside axon
    "blue",     # 5 periaxonal space
    "orange",   # 6 axon
    "purple",   # 7 myelin
])

def build_color_lut(cmap):
    lut = cmap(np.arange(len(cmap.colors)))[:, :3]
    return (lut * 255).astype(np.uint8)

def colorize(mask, lut):
    mask = mask.astype(np.int64)
    if mask.min() < 0 or mask.max() >= lut.shape[0]:
        raise ValueError("Mask labels exceed colormap range.")
    return lut[mask]

# Load a predicted mask (H, W), integer labels
pred_mask_path = "" # the path you want to save the result
mask = imread(pred_mask_path)
if mask.ndim == 3:
    mask = mask[..., 0]

lut = build_color_lut(cmap)
color_img = colorize(mask, lut)

imsave("pred_mask_color.png", color_img)

Notes

  • The choice of colors is purely for visualization and can be customized freely.

  • Users may ignore or recolor specific labels depending on their analysis goals.

  • Retaining all label IDs ensures compatibility with pipelines that combine outputs from multiple segmentation models.