API Reference

Auto-generated reference for the brain_mri_segmentation package via mkdocstrings. Each module below lists its public classes and functions with their docstrings and type signatures.

Package root

brain_mri_segmentation

Production-grade binary brain-tumor MRI segmentation (LGG / TCGA).

Data

datamodule

Lightning DataModule for binary brain-MRI segmentation.

dataset

Paired (image, mask) TIF dataset for binary segmentation.

prepare

Patient-level train/val/test split for the LGG Brain MRI dataset.

Models

factory

Model factory — returns a segmentation model by name.

lightning_module

Lightning module wrappers.

unet

Small U-Net (4 levels, 32->256 ch) for binary segmentation.

Training

train

Training entrypoint (Hydra-powered).

Evaluation

evaluate

Run model on test set, write reports/metrics.json (Dice + IoU + pixel accuracy).

Functions

Inference

predict

Inference CLI — load a checkpoint and predict a binary mask.

Functions

load_model

load_model(checkpoint_path: str | Path)

Load a Lightning module from checkpoint, rebuilding the backbone from hparams.

Source code in src/brain_mri_segmentation/inference/predict.py

def load_model(checkpoint_path: str | Path):
    """Load a Lightning module from checkpoint, rebuilding the backbone from hparams."""
    import torch

    from ..models import SegmentationModule, build_model

    ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
    hp = ckpt.get("hyper_parameters", {})
    model_name = hp.get("model_name")
    num_classes = hp.get("num_classes")
    if model_name is None or num_classes is None:
        raise ValueError(
            "Checkpoint missing model_name/num_classes hparams — re-train after upgrade."
        )
    backbone = build_model(model_name, num_classes=num_classes, pretrained=False)
    return SegmentationModule.load_from_checkpoint(str(checkpoint_path), model=backbone)

Serving

dependencies

Dependency injection — singleton model loader.

Functions

errors

Exception types and handlers.

main

FastAPI application.

routes

FastAPI routes.

schemas

Pydantic request/response schemas.

Utilities

hf_hub

HuggingFace Hub helpers.

logging

Structured logging configuration.

seed

Deterministic seeding across libraries.