Fracture Classification with Datamint and FracAtlas Dataset

This notebook demonstrates how to build an end-to-end binary classification pipeline using Datamint and the FracAtlas dataset.

Overview

You will learn how to:

• Set up a Datamint project for managing medical imaging data

• Upload the FracAtlas dataset to Datamint

• Train a model using the Datamint Trainer API — one line of code handles dataset loading, model creation, training, and experiment tracking

• Deploy the model for inference using Datamint’s model serving

Comparison with the Manual Workflow

The conventional way requires you to:

1. Define transforms, a Dataset subclass, and DataLoaders (~80 lines)

2. Define a full LightningModule with loss, metrics, and optimizer (~80 lines)

3. Configure MLflow logger, callbacks, and Trainer (~30 lines)

With the Trainer API, all of that is replaced by ~3 lines:

from datamint.lightning import EfficientNetV2Trainer

trainer = EfficientNetV2Trainer(project='FracAtlas')
results = trainer.fit()

Table of Contents

1. Setup: Create Project and Upload Dataset

2. Training with EfficientNetV2Trainer

3. Local Inference

4. Custom Model Architecture

5. Deployment

Required Dependencies

pip install datamint

Dataset Overview

The FracAtlas dataset contains 4,083 musculoskeletal radiographs annotated for fracture detection, localization, and segmentation. In this notebook, we focus on binary classification: determining whether an X-ray image shows a fracture or not.

from datamint import Api

PROJECT_NAME = "ClassificationTest"
api = Api()

1. Setup: Create Project and Load Dataset

datamint.examples.fracatlas_dataset downloads the FracAtlas dataset (This may take a few minutes), uploads the fractured / non-fractured X-ray images to a Datamint project, and creates a has_fracture classification annotation for each. The call is idempotent, if the project already exists, it’s returned as-is without re-uploading.

from datamint.examples import fracatlas_dataset

proj = fracatlas_dataset.create(project_name=PROJECT_NAME, api=api)
proj

1.1 Create Train/Validation/Test Splits

We use build_dataset to load the project and .split() to create reproducible splits. Calling .save() persists the assignments to the server — any later session can reload the same split with dataset.split() (no arguments).

Split	Percentage	Purpose
Train	80%	Model training
Validation	10%	Hyperparameter tuning, early stopping
Test	10%	Final model evaluation

from datamint.dataset import build_dataset

dataset = build_dataset(
    PROJECT_NAME,
    include_unannotated=False,
    image_categories_merge_strategy='mode',
    allow_external_annotations=True,
)

parts = dataset.split(train=0.8, val=0.1, test=0.1, seed=123)
parts.save(force=True)

for name, ds in parts.items():
    print(f'{name:6s}: {len(ds):4d} resources')

INFO:datamint.dataset.factory:Detected resource type 'image'; using ImageDataset.
INFO:datamint.dataset.base:Allowing external image label '('has_fracture', 'no')' not in project specs.
INFO:datamint.dataset.base:Allowing external image label '('has_fracture', 'yes')' not in project specs.

train : 3268 resources
val   :  408 resources
test  :  409 resources

2. Training with EfficientNetV2Trainer

This is where the Trainer API shines. Instead of manually defining:

• A Dataset subclass and DataLoaders

• A LightningModule with loss, metrics, and optimizer

• Callbacks, loggers, and a Lightning Trainer

• A deployment adapter

…you simply create an EfficientNetV2Trainer and call fit().

The trainer automatically:

1. Builds an ImageDataset and reads the saved project-scoped splits

2. Creates image augmentation pipelines (resize, flip, color jitter, normalize)

3. Instantiates EfficientNetV2-S pretrained on ImageNet at 384×384

4. Sets up CrossEntropy loss, Accuracy and F1 metrics

5. Configures MLflow logging, checkpointing, and early stopping

6. Trains, evaluates on the test split, and registers the best checkpoint in the MLflow Model Registry

from datamint.lightning import EfficientNetV2Trainer

trainer = EfficientNetV2Trainer(
    project=PROJECT_NAME,
    image_size=384,
    batch_size=8,
    max_epochs=3,
    accelerator='auto',
    pretrained=False,
    # model_name='MyModelName'  # defaults to project name
)

results = trainer.fit()

print("Test results:")
for metric_dict in results['test_results']:
    for k, v in metric_dict.items():
        print(f"  {k}: {v:.4f}")

# Open the Datamint project dashboard to view experiments and metrics
proj.show()

Inspect the Registered Model

Training automatically registered the model in Datamint’s model registry (backed by MLflow). Use api.models to inspect it without touching MLflow’s client directly:

registered_model = api.models.get_by_name(PROJECT_NAME)
latest_version = registered_model.get_latest_version()

print(f"Version: {latest_version.version}")
print(f"Task type: {latest_version.get_task_type()}")
print(f"Supported modes: {latest_version.get_supported_modes()}")
print("Metrics:")
for metric_name, value in latest_version.get_metrics().items():
    print(f"  {metric_name}: {value}")

3. Local Inference

The trainer automatically registers a deployment-ready model in MLflow after training. Load it for local prediction:

from datamint.mlflow import flavors as datamint_flavor

# Load the registered model
model_loaded = datamint_flavor.load_model(f'models:/{PROJECT_NAME}/latest')

# Predict on a resource from the dataset, and log the prediction back to Datamint
# as an annotation tagged source='model_deploy' (set log_predictions=True)
r = trainer.dataset[0]['resource']
prediction = model_loaded.predict([r], params={'log_predictions': True, 'model_name': PROJECT_NAME})
print(prediction)

[!TIP] You can also serve the model locally via REST API:

mlflow models serve -m "models:/{PROJECT_NAME}/latest" -p 5111 --env-manager virtualenv

4. Custom Model Architecture

The Trainer API lets you swap in a custom backbone while keeping everything else automatic.

There are two supported patterns:

1. (Recommended) Pass a ClassificationModule subclass class to model=. Datamint manages loss, metrics, MLflow logging, and deployment compatibility. Pass the class, not an instance.

2. Pass a fully constructed lightning.LightningModule instance to model= when you want complete control over training_step, validation_step, test_step, and configure_optimizers.

The example below uses ResNet-18 from torchvision as a custom backbone while keeping the full Datamint trainer workflow.

import torch
from torchvision.models import resnet18

from datamint.lightning import ImageClassificationTrainer
from datamint.lightning.trainers.lightning_modules import ClassificationModule


class FracAtlasClassifier(ClassificationModule):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Replace the default timm model with a custom backbone
        self.model = resnet18(weights='DEFAULT')
        self.model.fc = torch.nn.Linear(
            self.model.fc.in_features, self.hparams['num_classes']
        )

    def forward(self, x):
        return self.model(x)

    # OPTIONAL: override test_step for custom behaviour
    # def test_step(self, batch, batch_idx):
    #     images = batch['image']              # (B, C, H, W)
    #     labels = batch['image_categories']   # (B,)
    #     # ... custom logic ...
    #     return super().test_step(batch, batch_idx)


trainer_custom = ImageClassificationTrainer(
    project=PROJECT_NAME,
    model=FracAtlasClassifier,   # pass the class, not an instance
    image_size=480,
    batch_size=8,
    max_epochs=10,
    accelerator='auto',
)

results_custom = trainer_custom.fit()

print("Test results (custom model):")
for metric_dict in results_custom['test_results']:
    for k, v in metric_dict.items():
        print(f"  {k}: {v:.4f}")

4.1 Using a Fully Custom LightningModule

If you already have your own Lightning module, pass an instance. In this mode you own the training and evaluation logic:

import lightning as L
import torch
from torchvision.models import resnet50

class MyFullyCustomModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = resnet50(weights='DEFAULT')
        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 2)
        self.loss_fn = torch.nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images = batch['image']
        labels = batch['image_categories']
        loss = self.loss_fn(self(images), labels)
        self.log('train/loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        images = batch['image']
        labels = batch['image_categories']
        loss = self.loss_fn(self(images), labels)
        self.log('val/loss', loss)

    def test_step(self, batch, batch_idx):
        images = batch['image']
        labels = batch['image_categories']
        loss = self.loss_fn(self(images), labels)
        self.log('test/loss', loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-4)

trainer_full_custom = ImageClassificationTrainer(
    project=PROJECT_NAME,
    model=MyFullyCustomModule(),  # pass an instance
    image_size=480,
    batch_size=8,
    max_epochs=10,
    accelerator='auto',
)
results = trainer_full_custom.fit()

This still works with Datamint datasets, splits, MLflow logging, and checkpointing. The tradeoff is that a plain LightningModule does not automatically gain Datamint-native prediction methods for deployment. When you need Datamint-native serving, prefer ClassificationModule.

5. Deployment

The trainer automatically created and registered a deployment adapter in MLflow after training. Deploy it directly to the Datamint server:

job = api.deploy.start(
    model_name=PROJECT_NAME,
    model_alias="latest",
    with_gpu=False,
)

print("Deployment job started!")
print(f"Job ID: {job.id}")
print(f"Status: {job.status}")

job = api.deploy.get_by_id(job.id)

print(f"Job Status: {job.status}")
print(f"Progress: {job.progress_percentage}%")

if job.error_message:
    print(f"Error: {job.error_message}")

Model.is_deployed() reports the same thing, without tracking job IDs, once the job above reaches completed.

registered_model.is_deployed()

5.1 Remote Inference

r = api.resources.get_list(project_name=PROJECT_NAME, limit=1)[0]

inf_job = api.inference.submit(
    model_name=PROJECT_NAME,
    model_alias="latest",
    resource_id=r.id,
)
inf_job.wait()

print(inf_job.predictions)

proj.show()

Summary

Step	Description
Data Management	Uploaded 4,000+ images to Datamint with tags and annotations
Splits	Created 80/10/10 train/val/test splits with dataset.split().save()
Model Training	Fine-tuned EfficientNetV2-S with EfficientNetV2Trainer — one line
Experiment Tracking	Metrics and model automatically logged with MLflow
Custom Model	Swapped in ResNet-18 by subclassing ClassificationModule
Deployment	Model automatically registered and deployed to Datamint server

See Also

• Datamint Documentation

• FracAtlas Paper

• PyTorch Lightning Guide

• MLflow Model Registry