Inference

Inference#

MIPCandy provides a flexible inference system centered around the Predictor class, enabling prediction on various input formats including files, directories, tensors, and datasets.

Overview#

The inference module supports:

Flexible Input: Files, directories, tensors, sequences, or datasets
Lazy Model Loading: Models load only when first prediction is requested
Device Management: Automatic device placement and memory handling
Batch Processing: Efficient batch prediction with automatic padding
Easy Export: Save predictions to files with automatic naming

Quick Start#

from mipcandy_bundles.unet import UNetPredictor

# Create predictor from trained model
predictor = UNetPredictor(
    "experiments/UNet/20240901-1234",
    example_shape=(1, 128, 128),
    device="cuda"
)

# Predict single image
output = predictor.predict("path/to/image.nii.gz")

# Predict directory of images
outputs = predictor.predict("path/to/images/")

# Save predictions
predictor.predict_to_files("path/to/images/", "path/to/outputs/")

Creating a Predictor#

Basic Predictor Implementation#

To create a custom predictor, extend Predictor and implement build_network:

from typing import override
from torch import nn
from mipcandy.inference import Predictor
from mipcandy.types import AmbiguousShape

class MyPredictor(Predictor):
    @override
    def build_network(self, example_shape: AmbiguousShape) -> nn.Module:
        # Build network architecture based on example_shape
        model = MyNetwork(in_channels=example_shape[0])
        return model

# Usage
predictor = MyPredictor(
    "experiments/MyModel/20240901-1234",
    example_shape=(1, 128, 128),
    checkpoint="checkpoint_best.pth",
    device="cuda"
)

build_network is responsible only for constructing the model architecture. Checkpoint loading is handled automatically by the framework via safetensors – you do not need to call load_state_dict yourself.

Parameters:

experiment_folder: Path to trainer output directory
example_shape: A tuple describing the shape of a single input (e.g. (1, 128, 128) for single-channel 2D or (1, 128, 128, 128) for single-channel 3D). Used by build_network to configure the architecture.
checkpoint: Checkpoint filename (default: "checkpoint_best.pth")
device: Computing device (default: "cpu")

Lazy Model Loading#

Models are loaded only when first needed, saving memory when predictor is created but not immediately used:

# Predictor created but model not loaded yet
predictor = MyPredictor("experiments/model", (1, 128, 128), device="cuda")

# Model loads on first prediction
output = predictor.predict("image.nii.gz")  # Model loaded here

# Subsequent predictions reuse loaded model
output2 = predictor.predict("image2.nii.gz")  # Model already loaded

To explicitly load the model:

predictor = MyPredictor("experiments/model", (1, 128, 128), device="cuda")
predictor.lazy_load_model()  # Explicitly load model

Model Loading Flow#

Internally, when a prediction is first requested, lazy_load_model triggers the following chain:

build_network(example_shape) – constructs the model architecture
load_checkpoint(model, path) – loads weights from the safetensors checkpoint file
The model is moved to the configured device and set to eval mode

This separation means build_network should return a freshly initialized model without loading any weights.

Input Formats#

parse_predictant#

The parse_predictant function handles various input types:

from mipcandy.inference import parse_predictant
from mipcandy.data import Loader

# Single file
images, filenames = parse_predictant("image.nii.gz", Loader)
# images: list with 1 tensor
# filenames: ["image.nii.gz"]

# Directory
images, filenames = parse_predictant("images/", Loader)
# images: list with N tensors (one per file in directory)
# filenames: list of filenames

# Single tensor
tensor = torch.randn(1, 128, 128)
images, filenames = parse_predictant(tensor, Loader)
# images: [tensor]
# filenames: None

# List of files
images, filenames = parse_predictant(["img1.nii.gz", "img2.nii.gz"], Loader)
# images: list with 2 tensors
# filenames: ["img1.nii.gz", "img2.nii.gz"]

# List of tensors
images, filenames = parse_predictant([tensor1, tensor2], Loader)
# images: [tensor1, tensor2]
# filenames: None

Important

All elements in a sequence must have the same type (all strings or all tensors).

Parameters:

x: The input to parse (SupportedPredictant)
loader: The loader class to use for file loading (typically Loader)
as_label: Whether to load as label data (default: False)

Returns: tuple[list[torch.Tensor], list[str] | None] – a list of tensors and optionally a list of corresponding filenames.

Prediction Methods#

predict()#

Predict and return outputs as tensors:

predictor = MyPredictor("experiments/model", (1, 128, 128), device="cuda")

# Single image
output = predictor.predict("image.nii.gz")
# Returns: list[torch.Tensor] with 1 element

# Multiple images
outputs = predictor.predict("images_directory/")
# Returns: list[torch.Tensor] with N elements

# Tensors
tensor = torch.randn(1, 128, 128).cuda()
outputs = predictor.predict(tensor)
# Returns: list[torch.Tensor]

predict_image()#

Predict on a single tensor with optional batching:

# Single image (no batch dimension)
image = torch.randn(1, 128, 128).cuda()
output = predictor.predict_image(image, batch=False)
# Input shape: (C, H, W)
# Output shape: (C, H, W)

# Batch of images
images = torch.randn(4, 1, 128, 128).cuda()
outputs = predictor.predict_image(images, batch=True)
# Input shape: (B, C, H, W)
# Output shape: (B, C, H, W)

Parameters:

image: Input tensor (with or without batch dimension)
batch: Whether input has batch dimension (default: False)

When batch=False, the input is automatically unsqueezed before inference and squeezed back after. Padding and restoring modules (if configured) are applied transparently.

predict_to_files()#

Predict and save directly to files:

# Predict directory and save
filenames = predictor.predict_to_files(
    "input_images/",
    "output_predictions/"
)
# Saves predictions with original filenames
# Returns: list of filenames used, or None

# Custom filenames via save_predictions
outputs = predictor.predict("images/")
predictor.save_predictions(
    outputs,
    "output/",
    filenames=["pred_001.nii.gz", "pred_002.nii.gz"]
)

Returns: list[str] | None – the filenames used for saving, or None if inputs were tensors without associated filenames.

Callable Interface#

Predictors can be called directly:

predictor = MyPredictor("experiments/model", (1, 128, 128), device="cuda")

# Equivalent to predictor.predict()
outputs = predictor("images/")

Padding and Restoration#

Automatic Padding#

Predictors can optionally implement padding for inputs that don’t match required dimensions:

from typing import override
import torch
from torch import nn
from mipcandy.inference import Predictor
from mipcandy.common import Pad2d, Restore2d
from mipcandy.types import AmbiguousShape

class PaddedPredictor(Predictor):
    @override
    def build_network(self, example_shape: AmbiguousShape) -> nn.Module:
        return MyNetwork()

    @override
    def build_padding_module(self) -> nn.Module | None:
        # Pad to multiples of 128
        return Pad2d((128, 128))

    @override
    def build_restoring_module(self, padding_module: nn.Module | None) -> nn.Module | None:
        if padding_module:
            # Restore to original size
            return Restore2d(padding_module)
        return None

# Usage
predictor = PaddedPredictor("experiments/model", (1, 128, 128), device="cuda")

# Input: 100x100
# Automatically padded to 128x128
# Processed by model
# Automatically restored to 100x100
output = predictor.predict_image(torch.randn(1, 100, 100).cuda())

The padding and restoring modules are lazily loaded and cached for efficiency. Both modules are automatically moved to the predictor’s configured device.

The default implementations of build_padding_module and build_restoring_module return None, meaning no padding is applied unless explicitly overridden.

Dataset Integration#

Predictors work seamlessly with datasets:

from mipcandy.data import SimpleDataset, PathBasedUnsupervisedDataset

# Create dataset
dataset = SimpleDataset("test_images/", is_label=False, device="cuda")

# Predict entire dataset
outputs = predictor.predict(dataset)

# Process dataset case by case
for i, image in enumerate(dataset):
    output = predictor.predict_image(image)
    predictor.save_prediction(output, f"outputs/case_{i:03d}.nii.gz")

When a PathBasedUnsupervisedDataset is passed to predict or _predict, the predictor automatically extracts the file paths from the dataset for use as output filenames.

Saving Predictions#

save_prediction()#

Save a single prediction:

output = predictor.predict_image(image)
Predictor.save_prediction(output, "output.nii.gz")

This is a static method that delegates to save_image. It can be called on the class directly or on an instance.

save_predictions()#

Save multiple predictions with automatic or custom naming:

outputs = predictor.predict("images/")

# Automatic naming: prediction_00, prediction_01, ...
predictor.save_predictions(outputs, "output_folder/")

# Custom filenames
predictor.save_predictions(
    outputs,
    "output_folder/",
    filenames=["case1.nii.gz", "case2.nii.gz"]
)

Important

The output folder must already exist. save_predictions raises FileNotFoundError if the folder does not exist.

Automatic Naming Format: prediction_{i:0Nd} where N = ceil(ln(num_cases)) (natural logarithm).

The file extension is chosen automatically based on tensor dimensionality:

.png for 3D tensors with 1 or 3 channels (2D images)
.mha for all other shapes (3D volumes)

Example:

5 cases (N=2): prediction_00 to prediction_04
100 cases (N=5): prediction_00000 to prediction_00099
1000 cases (N=7): prediction_0000000 to prediction_0000999

Complete Example#

from typing import override
from os import PathLike
import torch
from torch import nn
from mipcandy.inference import Predictor
from mipcandy.types import AmbiguousShape, Device

class UNetPredictor(Predictor):
    def __init__(self, experiment_folder: str | PathLike[str], example_shape: AmbiguousShape, *,
                 checkpoint: str = "checkpoint_best.pth",
                 device: Device = "cuda") -> None:
        super().__init__(experiment_folder, example_shape, checkpoint=checkpoint, device=device)
        self.num_classes: int = 1

    @override
    def build_network(self, example_shape: AmbiguousShape) -> nn.Module:
        from my_models import UNet

        model = UNet(
            in_channels=example_shape[0],
            num_classes=self.num_classes
        )
        return model

# Inference pipeline
predictor = UNetPredictor(
    "experiments/UNet/20240901-1234",
    example_shape=(1, 128, 128),
    checkpoint="checkpoint_best.pth",
    device="cuda"
)

# Process test dataset
predictor.predict_to_files(
    "data/test_images/",
    "results/predictions/"
)

# Get predictions as tensors for further processing
outputs = predictor.predict("data/test_images/")
for i, output in enumerate(outputs):
    # Post-process predictions
    binary_mask = (output > 0.5).float()

    # Save processed result
    predictor.save_prediction(binary_mask, f"results/binary/case_{i:03d}.nii.gz")

Inference

Contents

Inference#

Overview#

Quick Start#

Creating a Predictor#

Basic Predictor Implementation#

Lazy Model Loading#

Model Loading Flow#

Input Formats#

parse_predictant#

Prediction Methods#

predict()#

predict_image()#

predict_to_files()#

Callable Interface#

Padding and Restoration#

Automatic Padding#

Dataset Integration#

Saving Predictions#

save_prediction()#

save_predictions()#

Complete Example#