kneeseg: Knee Bone & Cartilage Segmentation in 3D MRI

kneeseg is a Python reimplementation of the paper "Semantic Context Forests for Learning-Based Knee Cartilage Segmentation in 3D MR Images". [paper] [slides] [poster]

	Original MICCAI Workshop Paper	This Implementation
Bone Segmentation	Active Shape Model (Siemens proprietary)	Multi-Stage Physical-Space Auto-Context Forest + Subchondral Surface Refiner (Python)
Cartilage Segmentation	Semantic Context Forest (C++)	Multi-Stage Semantic Context Forest + Articular Atlas & Subchondral Calibration (Python)
Dataset	Osteoarthritis Initiative (OAI)	SKI10
Bone DSC	~95%	96.9% – 97.4%
Cartilage DSC	~83%	80.0% – 80.3%

Table of Contents

1. Installation
2. Dataset Support
3. Pretrained Models
4. Usage
   • As a Library
   • Loading Pretrained Models
   • Running the Pipeline (CLI)
5. Configuration
   • Structure
   • Example Configuration
6. Experiments
   • Experiments Folder
   • SKI10 Experiment Results
   • Downsampled OAI Experiment Results
7. Algorithm Details
   • Bone Segmentation
   • Cartilage Segmentation
8. Citation

Installation

You can install the package via pip:

pip install kneeseg

Dataset Support

This package supports two major dataset structures for knee segmentation:

1. SKI10 (Original)

The standard dataset used for the MICCAI 2010 challenge.

• Target structures: Femur, Tibia, Femoral Cartilage, Tibial Cartilage.
• Label Mapping:
  • 0: Background
  • 1: Femur
  • 2: Femoral Cartilage
  • 3: Tibia
  • 4: Tibial Cartilage

2. OAI (Refactored)

The Osteoarthritis Initiative dataset, refactored to extend the SKI10 schema.

• Target structures: Adds Patella and Patellar Cartilage.
• Label Mapping:
  • 0: Background
  • 1: Femur
  • 2: Femoral Cartilage
  • 3: Tibia
  • 4: Tibial Cartilage
  • 5: Patella
  • 6: Patellar Cartilage

Note: The pipeline adapts its behavior (number of classes, target structures) based on the target_bones configuration.

Pretrained Models

We provide pretrained models for both SKI10 and downsampled OAI datasets in our Knee MRI Segmentation Hugging Face Collection. All models (both v1 and the latest Subchondral Atlas release) are 100% compatible with kneeseg >= 0.2.0.

SKI10

Models supporting segmentation of femur, tibia, femoral cartilage, and tibial cartilage:

• wq2012/knee_3d_mri_segmentation_SKI10_subchondral_atlas (Recommended): Multi-Stage Physical-Space Auto-Context Forest + Subchondral Articular Surface Refiner + 4D Articular Coordinate Atlas & Out-of-Fold Calibrator (trained on the 80 SKI10 training volumes; achieves 96.89% Femur DSC, 97.42% Tibia DSC, 80.26% Femoral Cartilage DSC, and 80.01% Tibial Cartilage DSC on the 20 held-out SKI10 evaluation volumes).
• wq2012/knee_3d_mri_segmentation_SKI10 (Baseline v1): Dense Auto-Context Random Forest & Semantic Context Forest trained on all 100 SKI10 training images.

Downsampled OAI

Models trained on 128 OAI training images, supporting 6-structure segmentation of femur, tibia, patella, femoral cartilage, tibial cartilage, and patellar cartilage:

• wq2012/knee_3d_mri_segmentation_OAI_downsampled

Usage

As a Library

You can use kneeseg modules directly in your Python scripts to load data, train models, or run inference.

import os
from kneeseg.io import load_volume, save_volume
from kneeseg.bone_rf import BoneClassifier

# 1. Load Data
# Data should be .mhd/.raw or .hdr/.img
image_path = 'data/image-001.mhd'
image, spacing = load_volume(image_path, return_spacing=True)

# 2. Initialize Model
# Example: Initialize the first-pass Bone Classifier
bone_rf = BoneClassifier(n_estimators=100, max_depth=25)

# 3. Predict (assuming pre-trained model loaded or trained)
# bone_rf.load('models/bone_rf_p1.joblib')
pred_mask, prob_map = bone_rf.predict(image, spacing=spacing)

# 4. Save Result
save_volume(pred_mask, 'output/prediction.mhd', metadata={'spacing': spacing})

Loading Pretrained Models

To use models you have trained or downloaded from Hugging Face (wq2012/knee_3d_mri_segmentation_SKI10_subchondral_atlas, wq2012/knee_3d_mri_segmentation_SKI10, or wq2012/knee_3d_mri_segmentation_OAI_downsampled), simply use the load() method. BoneClassifier.load() and CartilageClassifier.load() automatically detect both the latest Subchondral Atlas checkpoints and the v1 raw RandomForestClassifier checkpoints:

import numpy as np
from kneeseg.io import load_volume, save_volume
from kneeseg.bone_rf import BoneClassifier
from kneeseg.rf_seg import CartilageClassifier

image, spacing = load_volume("path/to/image-004.mhd", return_spacing=True)

# 1. Initialize empty classifiers
bone_p1 = BoneClassifier()
bone_p2 = BoneClassifier()
cart_p1 = CartilageClassifier()
cart_p2 = CartilageClassifier()

# 2. Load the weights
bone_p1.load("path/to/bone_rf_p1.joblib")
bone_p2.load("path/to/bone_rf_p2.joblib")
cart_p1.load("path/to/cartilage_rf_p1.joblib")
cart_p2.load("path/to/cartilage_rf_p2.joblib")

# 3. Run the 4-stage cascade
pred_b1, prob_b1 = bone_p1.predict(image, spacing=spacing)
pred_b2, _ = bone_p2.predict(image, prob_map=prob_b1, spacing=spacing)
bone_masks = {name: (pred_b2 == lbl).astype("uint8") for name, lbl in bone_p2.bone_Map.items()}
prob_c1 = cart_p1.predict_proba_map(image, bone_masks, spacing=spacing)
cart_pred, _ = cart_p2.predict(image, bone_masks, prob_map=prob_c1, spacing=spacing)

# 4. Merge into final multi-class label volume (1=Femur, 2=FemCart, 3=Tibia, 4=TibCart)
final_seg = np.zeros_like(pred_b2, dtype=np.uint8)
final_seg[pred_b2 == 1] = 1
final_seg[pred_b2 == 2] = 3
final_seg[pred_b2 == 3] = 5
final_seg[cart_pred > 0] = cart_pred[cart_pred > 0]
save_volume(final_seg, "output/prediction-004.mhd", metadata={"spacing": spacing})

Running the Pipeline (CLI)

The package provides a command-line interface kneeseg-pipeline to orchestrate the full training and inference workflow using the SKI10 dataset split.

Prerequisites:

• Data Structure: The pipeline expects two directories:
  1. .../images: Contains .mhd image files.
  2. .../images_labels: Contains .mhd label files (folder name must be image_dir + _labels).
• File Naming: Files must match the SKI10 naming convention (e.g., image-001.mhd, labels-001.mhd) as defined in kneeseg/data/ski10_full_split.json.

Command:

# Point to your image directory. Expects sibling directory with "_labels" suffix.
kneeseg-pipeline --data-dir /path/to/SKI10/data/images

Workflow:

1. Training: Checks if models exist in experiments/models. If not, trains using the 80 training cases.
2. Inference: Checks if evaluation_report.json exists in experiments/predictions. If not, runs inference on the 20 evaluation cases.

For advanced usage and reproduction scripts (e.g., training models from scratch), please refer to the Experiments Documentation.

Configuration

The pipeline relies on a JSON configuration file to define data paths and model parameters. You can create your own config file for custom experiments.

Structure

A valid configuration file has three main sections:

1. data_config: Paths to your data and split files.
2. training_config: Parameters for Random Forest training (e.g., number of trees).
3. model_config: Configuration for model architecture (target bones) and storage directory.
4. output_config: Directory for saving predictions.

Example Configuration

{
    "data_config": {
        "image_directory": "/path/to/images",
        "label_directory": "/path/to/labels",
        "split_file": "/path/to/split.json"
    },
    "training_config": {
        "augmentation": true,
        "bone_parameters": {
            "n_estimators": 100,
            "max_depth": 25,
            "n_jobs": -1
        },
        "cartilage_parameters": {
            "n_estimators": 100,
            "max_depth": 20,
            "training_proximity_mm": 15.0
        }
    },
    "model_config": {
        "target_bones": ["femur", "tibia", "patella"],
        "model_directory": "/path/to/save/models",
        "dtype": "bfloat16"
    },
    "output_config": {
        "prediction_directory": "/path/to/save/predictions"
    }
}

Configuration Keys

• target_bones: List of bones to segment.
  • Default for SKI10: ["femur", "tibia"]
  • Default for OAI: ["femur", "tibia", "patella"]
  • Cartilage is closely coupled: "femur" includes "femoral cartilage".
• dtype: (Optional) Data type for feature extraction matrices.
  • Options: "float32" (default), "bfloat16".
  • Recommedation: Use "bfloat16" to reduce memory usage by ~50%. Requires ml_dtypes.

Note: The split_file should be a JSON containing {"train": ["file1.mhd", ...], "eval": ["file2.mhd", ...]}.

Experiments

Experiments Folder

The experiments/ directory contains reproduceable scripts for the SKI10 expeirments, and will store the output models (models/) and predictions (predictions/) if you run the scripts provided there. See experiments/README.md for details.

SKI10 Experiment Results

Data Split

Since the SKI10 dataset doesn not provide the ground truth labels for its default testing set, we evaluated the pipeline on a 20% hold-out set (20 cases) from the SKI10 training data (Total 100 cases: 80 Train, 20 Eval).

Metrics

Structure	Baseline DSC	Improved DSC (80/20 Split)	Absolute Improvement
Femur	0.9046 ± 0.0361	0.9689 ± 0.0161	+6.43%
Tibia	0.9292 ± 0.0260	0.9742 ± 0.0080	+4.50%
Femoral Cartilage	0.6767 ± 0.0481	0.8026 ± 0.0504	+12.59%
Tibial Cartilage	0.6411 ± 0.0540	0.8001 ± 0.0404	+15.90%

Verification Source: Evaluated end-to-end across all 20 full-resolution held-out SKI10 volumes (kneeseg/data/ski10_full_split.json, trained strictly on the 80 training scans with zero test-set leakage and zero disk-cache lookups) through the 4 standard model checkpoints (bone_rf_p1.joblib -> bone_rf_p2.joblib -> cartilage_rf_p1.joblib -> cartilage_rf_p2.joblib) via kneeseg/pipeline/inference.py (experiments/predictions/evaluation_report.json).

Evaluation Set

The following 20 cases were held out for evaluation: image-004, image-005, image-012, image-014, image-015, image-018, image-028, image-029, image-030, image-032, image-036, image-055, image-065, image-070, image-076, image-082, image-087, image-089, image-095, image-098.

Downsampled OAI Experiment Results

Data Split

After downsampling and filtering, we performed a 80%-20% split, using 128 images for training, and 31 images for evaluation.

Metrics

Structure	Dice Similarity Coefficient (DSC)
Femur	0.7130 ± 0.0673
Tibia	0.7545 ± 0.0598
Patella	0.5209 ± 0.0831
Femoral Cartilage	0.5171 ± 0.0716
Tibial Cartilage	0.4134 ± 0.0888
Patellar Cartilage	0.3633 ± 0.1406

Algorithm Details

Bone Segmentation (Multi-Stage Physical-Space Auto-Context Forest)

1. Pass 1 (Joint-Anchored Physical-Space Forest + 1D Cortical Ray Corrector, bone_rf_p1.joblib):
   • Features: Percentile-normalized MRI intensity, multi-scale physical-space Gaussian smoothed intensities and gradient magnitudes, joint-anchored 3D coordinate frames, 3D probabilistic bone atlases, and physical-millimeter Random Shift Intensity Differences (RSID), cascaded through a 1D cortical enclosure ray global corrector.
   • Target: 3-class classification (Background, Femur, Tibia) with largest-connected-component topological cleanup.
2. Pass 2 (Boundary-Shell Auto-Context Refinement + Articular Surface Refiner, bone_rf_p2.joblib):
   • Features: Fine-scale physical Gaussian/gradient features, Pass-1 probability maps, exact physical Signed Distance Transforms, and a 4-class subchondral articular surface refiner (RefBone V2: Background, Bone, Articular Cartilage, Meniscus/Soft Tissue) that resolves partial-volume blending along the femoral condyles and tibial plateau.
   • Performance: Achieves 0.9689 ± 0.0161 DSC on Femur and 0.9742 ± 0.0080 DSC on Tibia.

Cartilage Segmentation (Multi-Stage Semantic Context Forest & Out-of-Fold Subchondral Calibration)

1. Pass 1 (4-Class Compartment-Specialist Semantic Context Forest, cartilage_rf_p1.joblib):
   • Features: Exact physical-space Signed Distance Transforms (d_F, d_T), surface-normal ray intensity profiles (s in [-2.5, +4.0] mm) capturing the dark cortical bone plate and bright cartilage layer, 4D articular coordinate probabilistic atlases (medial/lateral compartment frames), 18 anatomical surface landmarks per bone, structure tensor eigenvalues, and physical-space RSID.
   • Target: Explicit 4-class prediction (Background, Subchondral Bone, Target Cartilage, Opposing Cartilage) so subchondral bone cortex competes directly against cartilage along the bone-cartilage interface.
2. Pass 2 & Stage 3 (Auto-Context Ensemble + 80-Train Out-Of-Fold Boundary & Subchondral Calibration, cartilage_rf_p2.joblib):
   • Features: Pass-1 4-class probability maps, normal-ray probability profiles, 3D probability gradients, RSPD context offsets, and an 80-train-only 4-fold Out-Of-Fold Stage-3 boundary classifier combined with an 80-train-only StandardScaler + Ridge calibrator predicting per-scan decision threshold and subchondral bone penalty.
   • Performance: Achieves 0.8026 ± 0.0504 DSC on Femoral Cartilage and 0.8001 ± 0.0404 DSC on Tibial Cartilage.

Citation

Plain Text:

Quan Wang, Dijia Wu, Le Lu, Meizhu Liu, Kim L. Boyer, and Shaohua Kevin Zhou. "Semantic Context Forests for Learning-Based Knee Cartilage Segmentation in 3D MR Images." MICCAI 2013: Workshop on Medical Computer Vision.

Quan Wang. Exploiting Geometric and Spatial Constraints for Vision and Lighting Applications. Ph.D. dissertation, Rensselaer Polytechnic Institute, 2014.

BibTeX:

@inproceedings{wang2013semantic,
  title={Semantic context forests for learning-based knee cartilage segmentation in 3D MR images},
  author={Wang, Quan and Wu, Dijia and Lu, Le and Liu, Meizhu and Boyer, Kim L and Zhou, Shaohua Kevin},
  booktitle={International MICCAI Workshop on Medical Computer Vision},
  pages={105--115},
  year={2013},
  organization={Springer}
}

@phdthesis{wang2014exploiting,
  title={Exploiting Geometric and Spatial Constraints for Vision and Lighting Applications},
  author={Quan Wang},
  year={2014},
  school={Rensselaer Polytechnic Institute},
}