Deploying an Externally Trained Model with Datamint

This notebook shows how to take a model trained outside of Datamint — in your own script, a Jupyter notebook, or any third-party framework — and register, deploy, and run inference with it using Datamint’s serving infrastructure.

See Bringing an External Model into Datamint for a narrative walkthrough of this same flow (integrate, log, deploy).

When to use this tutorial

Use this approach when you:

  • Already have a trained PyTorch model (checkpoint, state_dict, or full nn.Module)

  • Trained with a custom loop, Keras, Hugging Face, or any other framework

  • Want to deploy a third-party pre-trained model (foundation model, off-the-shelf architecture)

Comparison with the Trainer API

Step

Trainer API

This tutorial

Training

Handled automatically

manual

MLflow logging

Handled automatically

datamint.mlflow.flavors.log_model(...)

Adapter code

None required

Subclass DatamintModel

Deployment

Handled automatically

api.deploy.start(...)

What You’ll Learn

  1. Wrap a plain nn.Module in a DatamintModel adapter

  2. Implement predict_default to produce Datamint annotations

  3. Log and register the model in MLflow

  4. Test inference locally

  5. Deploy the model and run remote inference

Required Dependencies

pip install datamint segmentation-models-pytorch albumentations

0. Setup

[ ]:

from datamint import Api

PROJECT_NAME = "ExternalModel_Tutorial"
MODEL_NAME = PROJECT_NAME  # MLflow registered model name

api = Api()
proj = api.projects.create(
    name=PROJECT_NAME,
    description="Tutorial: deploying an externally trained segmentation model",
    exists_ok=True,
)

1. Train or Load Your Model

In a real scenario you already have a checkpoint. Here we train a minimal UNet++ on a toy dataset so the notebook runs end-to-end without external data.

Replace the cell below with your own model loading code, for example:

import torch
import segmentation_models_pytorch as smp

model = smp.UnetPlusPlus(encoder_name='resnet34', in_channels=3, classes=1)
state = torch.load('my_checkpoint.pth', map_location='cpu')
model.load_state_dict(state)
model.eval()

EXAMPLE TRAINING CODE:

[ ]:

import segmentation_models_pytorch as smp
import torch

CLASS_NAMES = ['lesion']  # foreground class names; one per output channel
IMAGE_SIZE = 256

# ── Build architecture ──────────────────────────────────────────────────────
net = smp.UnetPlusPlus(
    encoder_name='resnet34',
    encoder_weights='imagenet',
    in_channels=3,
    classes=len(CLASS_NAMES),
)

# ── (Optional) quick training on a toy batch ────────────────────────────────
# Replace this block with your real training loop / checkpoint loading.
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
loss_fn = smp.losses.DiceLoss(mode='binary')

net.train()
for _ in range(3):  # 3 fake gradient steps
    x = torch.randn(2, 3, IMAGE_SIZE, IMAGE_SIZE)
    y = (torch.rand(2, len(CLASS_NAMES), IMAGE_SIZE, IMAGE_SIZE) > 0.5).float()
    loss = loss_fn(net(x), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

net.eval()
print(f"Model ready — {sum(p.numel() for p in net.parameters()):,} parameters")

2. Create a DatamintModel Adapter

Datamint’s deployment server calls your predict_default method with a list of Resource objects and expects a list[list[Annotation]] back — one annotation list per resource.

You wrap your model by subclassing DatamintModel and implementing predict_default:

DatamintModel
└── predict_default(model_input: list[Resource], **kwargs) → list[list[Annotation]]

Key helpers available inside predict_default

Helper

Returns

self.inference_device

'cuda' or 'cpu' (set by the server)

self.get_pytorch_model()

The nn.Module you passed as torch_model

resource.fetch_file_data(auto_convert=True)

PIL Image / DICOM etc.

Supported annotation types

Import from datamint.entities.annotations:

  • ImageSegmentation — 2-D binary mask

  • CategoryAnnotation — class label + optional probability

  • BoxAnnotation — box from two corner points

[ ]:

import albumentations as A
import cv2
import numpy as np
from albumentations.pytorch import ToTensorV2

from datamint.entities.annotations import ImageSegmentation
from datamint.mlflow.flavors.model import DatamintModel, ModelSettings
from datamint.mlflow.flavors.task_type import TaskType


class SegmentationAdapter(DatamintModel):
    """Wraps a plain nn.Module for Datamint segmentation inference."""

    task_type = TaskType.IMAGE_SEGMENTATION

    def __init__(
        self,
        torch_model: torch.nn.Module,
        class_names: list[str],
        image_size: int = 256,
        threshold: float = 0.5,
        need_gpu: bool = False,
    ) -> None:
        super().__init__(
            torch_model=torch_model,
            settings=ModelSettings(need_gpu=need_gpu),
        )
        self.class_names = class_names
        self.image_size = image_size
        self.threshold = threshold

        # Preprocessing applied to every image at inference time
        self._transform = A.Compose([
            A.Resize(image_size, image_size),
            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            ToTensorV2(),
        ])

    def predict_default(
        self,
        model_input: list,  # list[Resource]
        **kwargs,
    ) -> list[list[ImageSegmentation]]:
        device = self.inference_device
        model = self.get_pytorch_model().to(device).eval()

        results = []
        for resource in model_input:
            # ── Load image ──────────────────────────────────────────────────
            img = np.array(resource.fetch_file_data(auto_convert=True, use_cache=True))
            if img.ndim == 2:                         # grayscale → 3-channel
                img = np.stack([img, img, img], axis=-1)
            elif img.shape[-1] == 4:                  # RGBA → RGB
                img = img[..., :3]
            orig_h, orig_w = img.shape[:2]

            # ── Preprocess ──────────────────────────────────────────────────
            tensor = self._transform(image=img)['image'].unsqueeze(0).to(device)  # (1, 3, H, W)

            # ── Forward pass ────────────────────────────────────────────────
            with torch.inference_mode():
                logits = model(tensor)                            # (1, C, H, W)
            probs = logits.sigmoid().squeeze(0).cpu().numpy()    # (C, H, W)  [0, 1]

            # ── Build annotations ───────────────────────────────────────────
            annotations = [
                ImageSegmentation(
                    name=self.class_names[i],
                    segmentation_data=cv2.resize(
                        (probs[i] > self.threshold).astype(np.uint8),
                        (orig_w, orig_h),
                        interpolation=cv2.INTER_NEAREST,
                    ),
                )
                for i in range(len(self.class_names))
            ]
            results.append(annotations)

        return results


# Instantiate the adapter with the trained model
adapter = SegmentationAdapter(
    torch_model=net,
    class_names=CLASS_NAMES,
    image_size=IMAGE_SIZE,
)
print("Adapter ready.")

2.1 Smoke-test the adapter locally

Before logging to MLflow it’s worth making sure predict_default runs without errors using a fake resource that just wraps a local file.

[ ]:

import io

from PIL import Image

from datamint.entities.resource import LocalResource

# Create a dummy in-memory PNG as a LocalResource
buf = io.BytesIO()

img = Image.fromarray(np.random.randint(0, 255, (300, 400, 3), dtype=np.uint8))
img.save(buf, format='PNG')
dummy_resource = LocalResource(raw_data=buf.getvalue(), filename='dummy.png')

predictions = adapter.predict([dummy_resource])

print(f"Received {len(predictions)} result(s)")
print(f"Annotations per resource: {[len(p) for p in predictions]}")
for ann in predictions[0]:
    print(f"  {ann.name!r}  mask shape={ann.mask.shape}  dtype={ann.mask.dtype}")

[ ]:

# upload resource to Datamint (optional, but simulates real usage more closely)
import tempfile

with tempfile.NamedTemporaryFile(suffix='.png') as tmp:
    img.save(tmp.name)
    api.resources.upload_resource(tmp.name, tags=['dummy'],
                                  publish_to=proj)

3. Log and Register the Model in MLflow

datamint.mlflow.flavors.log_model is the single call that:

  1. Serialises your adapter (including the embedded nn.Module) as an MLflow artifact

  2. Records the task_type so the Datamint server knows how to display predictions

  3. Pins the datamint and medimgkit package versions in requirements.txt

  4. (Optionally) registers the model in the MLflow Model Registry under a name you choose

Why call datamint.mlflow.set_project first?

set_project configures the MLflow tracking URI to point to your Datamint server and associates the run with the correct project, so metrics and artifacts are visible in the Datamint dashboard.

[ ]:

import mlflow

import datamint.mlflow as datamint_mlflow
from datamint.mlflow.flavors import log_model

# Point MLflow at the Datamint server and associate runs with the project
datamint_mlflow.set_project(PROJECT_NAME)
mlflow.set_experiment(PROJECT_NAME) # Can be any name; doesn't have to match the project name

with mlflow.start_run(run_name='external_model_upload') as run:
    # Log any metadata you want to track alongside the model
    mlflow.log_params({
        'encoder': 'resnet34',
        'image_size': IMAGE_SIZE,
        'class_names': str(CLASS_NAMES),
        'framework': 'segmentation_models_pytorch',
    })

    model_info = log_model(
        adapter,
        task_type=TaskType.IMAGE_SEGMENTATION,
        name='segmentation_model',          # artifact sub-path inside the run
        registered_model_name=MODEL_NAME,   # registers in the Model Registry
    )

print(f"Model URI : {model_info.model_uri}")
print(f"Run ID    : {run.info.run_id}")

3.1 Assign the champion alias

The deployment API (api.deploy.start) resolves models by alias. Set the champion (or any of your choice) alias on the version we just registered.

[ ]:

from mlflow import MlflowClient

client = MlflowClient()

# The newly registered version is always the highest version number
versions = client.search_model_versions(f"name='{MODEL_NAME}'")
latest_version = max(versions, key=lambda v: int(v.version))

client.set_registered_model_alias(MODEL_NAME, 'champion', latest_version.version)
print(f"Alias 'champion' → version {latest_version.version}")

4. Verify: Load and Run Local Inference

Before deploying, load the registered model back from MLflow and verify that predict works correctly end-to-end.

[ ]:

from datamint.mlflow import flavors as datamint_flavor

model = api.models.get_by_name(MODEL_NAME)
loaded_model = model.get_latest_version(alias='champion').load_model()

print(f"Loaded model type : {type(loaded_model).__name__}")
print(f"Task type         : {loaded_model.task_type}")
print(f"Inference device  : {loaded_model.inference_device}")

[ ]:

# Run inference on a real resource from the project
resources = list(api.resources.get_list(project_name=PROJECT_NAME, limit=3))

if resources:
    predictions = loaded_model.predict(resources[:1])
    for ann in predictions[0]:
        print(f"  class={ann.name!r}  mask={ann.mask.shape}  positive_pixels={ann.mask.sum()}")
else:
    # Fall back to the dummy resource from section 2.1
    predictions = loaded_model.predict([dummy_resource])
    print("No resources in project — tested on dummy image.")
    for ann in predictions[0]:
        print(f"  class={ann.name!r}  mask={ann.mask.shape}")

5. Deploy to the Datamint Server

[ ]:

job = api.deploy.start(
    model_name=MODEL_NAME,
    model_alias='champion',
    with_gpu=False,
)

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

[ ]:

import time

# Poll until the deployment is complete (usually 5–10 minutes)
while True:
    job = api.deploy.get_by_id(job.id)
    print(f"Status: {job.status}  ({job.progress_percentage:.0f}%)")
    if job.status in ('completed', 'failed', 'cancelled'):
        break
    time.sleep(15)

if job.status == 'completed':
    print("Deployment successful!")
else:
    print(f"Deployment failed: {job.error_message}")
    print(job.build_logs)

6. Remote Inference

Once deployed, submit resources for server-side inference.

[ ]:

# Pick a resource to run inference on
r = api.resources.get_list(project_name=PROJECT_NAME, limit=1)[0]

inf_job = api.inference.submit(
    model_name=MODEL_NAME,
    model_alias='champion',
    resource_id=r.id,
)
inf_job.wait()  # blocks until inference is complete

print(f"Inference complete — {len(inf_job.predictions)} result(s)")

[ ]:

from matplotlib import pyplot as plt

preds = inf_job.predictions[0]

fig, axes = plt.subplots(1, len(preds), figsize=(5 * max(len(preds), 1), 5))
if len(preds) == 1:
    axes = [axes]

for ax, ann in zip(axes, preds):
    ax.imshow(ann.mask, cmap='gray')
    ax.set_title(ann.name)
    ax.axis('off')

plt.suptitle(f"Remote inference on: {r.filename}")
plt.tight_layout()
plt.show()

7. Updating the Model (New Version)

When you retrain or fine-tune your model, log a new version and move the champion alias:

# Re-train / fine-tune your net ...

adapter_v2 = SegmentationAdapter(torch_model=net_v2, class_names=CLASS_NAMES)

datamint_mlflow.set_project(PROJECT_NAME)
mlflow.set_experiment(PROJECT_NAME)

with mlflow.start_run(run_name='external_model_v2'):
    mlflow.log_param('version', 2)
    model_info = log_model(
        adapter_v2,
        task_type=TaskType.IMAGE_SEGMENTATION,
        name='segmentation_model',
        registered_model_name=MODEL_NAME,
    )

# Advance the alias so deployments pick up the new version automatically
versions = client.search_model_versions(f"name='{MODEL_NAME}'")
latest_version = max(versions, key=lambda v: int(v.version))
client.set_registered_model_alias(MODEL_NAME, 'champion', latest_version.version)

# Re-deploy
api.deploy.start(model_name=MODEL_NAME, model_alias='champion', with_gpu=False)

8. Custom Prediction Modes (Advanced)

For large images you may want slice-by-slice or frame-by-frame inference. Add extra prediction hooks using @prediction_mode:

from datamint.mlflow.flavors.prediction_router import prediction_mode
from datamint.mlflow.flavors.prediction_modes import PredictionMode

class MyAdapter(DatamintModel):
    task_type = TaskType.IMAGE_SEGMENTATION

    def predict_default(self, model_input, **kwargs):
        ...  # full-image inference

    @prediction_mode(PredictionMode.SLICE)
    def predict_slice(self, model_input, slice_idx, **kwargs):
        ...  # called when mode='slice' is passed as a parameter

Supported modes are defined in datamint.mlflow.flavors.prediction_modes.PredictionMode.

Summary

To deploy any externally trained model with Datamint:

Step

Code

  1. Wrap your model

Subclass DatamintModel, implement predict_default

  1. Configure MLflow

datamint.mlflow.set_project(project_name)

  1. Log & register

log_model(adapter, task_type=..., registered_model_name=...)

  1. Set alias

client.set_registered_model_alias(name, 'latest', version)

  1. Deploy

api.deploy.start(model_name=..., model_alias='latest')

  1. Infer

api.inference.submit(model_name=..., resource_id=...)

The only part unique to your model is the adapter class — everything else is identical regardless of architecture, framework, or task type.