{ "cells": [ { "cell_type": "markdown", "id": "5794f1e6", "metadata": {}, "source": [ "# Deploying an Externally Trained Model with Datamint\n", "\n", "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.\n", "\n", "> See [Bringing an External Model into Datamint](https://sonanceai.github.io/datamint-python-api/inference.html) for a narrative walkthrough of this same flow (integrate, log, deploy).\n", "\n", "## When to use this tutorial\n", "\n", "Use this approach when you:\n", "- Already have a trained PyTorch model (checkpoint, `state_dict`, or full `nn.Module`)\n", "- Trained with a custom loop, Keras, Hugging Face, or any other framework\n", "- Want to deploy a third-party pre-trained model (foundation model, off-the-shelf architecture)\n", "\n", "## Comparison with the Trainer API\n", "\n", "| Step | [Trainer API](segmentation_2d_trainer_BUSI_tutorial.ipynb) | This tutorial |\n", "|------|-----------------------------------------------------------|---------------|\n", "| Training | Handled automatically | manual |\n", "| MLflow logging | Handled automatically | `datamint.mlflow.flavors.log_model(...)` |\n", "| Adapter code | None required | Subclass `DatamintModel` |\n", "| Deployment | Handled automatically | `api.deploy.start(...)` |\n", "\n", "## What You'll Learn\n", "\n", "1. Wrap a plain `nn.Module` in a `DatamintModel` adapter\n", "2. Implement `predict_default` to produce Datamint annotations\n", "3. Log and register the model in MLflow\n", "4. Test inference locally\n", "5. Deploy the model and run remote inference\n", "\n", "## Required Dependencies\n", "\n", "```bash\n", "pip install datamint segmentation-models-pytorch albumentations\n", "```" ] }, { "cell_type": "markdown", "id": "7a9756c6", "metadata": {}, "source": [ "## 0. Setup" ] }, { "cell_type": "code", "execution_count": null, "id": "e8b17edf", "metadata": {}, "outputs": [], "source": [ "from datamint import Api\n", "\n", "PROJECT_NAME = \"ExternalModel_Tutorial\"\n", "MODEL_NAME = PROJECT_NAME # MLflow registered model name\n", "\n", "api = Api()\n", "proj = api.projects.create(\n", " name=PROJECT_NAME,\n", " description=\"Tutorial: deploying an externally trained segmentation model\",\n", " exists_ok=True,\n", ")" ] }, { "cell_type": "markdown", "id": "69ec0bb2", "metadata": {}, "source": [ "## 1. Train or Load Your Model\n", "\n", "In a real scenario you already have a checkpoint. Here we train a minimal UNet++ on a toy dataset\n", "so the notebook runs end-to-end without external data.\n", "\n", "Replace the cell below with your own model loading code, for example:\n", "\n", "```python\n", "import torch\n", "import segmentation_models_pytorch as smp\n", "\n", "model = smp.UnetPlusPlus(encoder_name='resnet34', in_channels=3, classes=1)\n", "state = torch.load('my_checkpoint.pth', map_location='cpu')\n", "model.load_state_dict(state)\n", "model.eval()\n", "```\n", "\n", "EXAMPLE TRAINING CODE:" ] }, { "cell_type": "code", "execution_count": null, "id": "79833da1", "metadata": {}, "outputs": [], "source": [ "import segmentation_models_pytorch as smp\n", "import torch\n", "\n", "CLASS_NAMES = ['lesion'] # foreground class names; one per output channel\n", "IMAGE_SIZE = 256\n", "\n", "# ── Build architecture ──────────────────────────────────────────────────────\n", "net = smp.UnetPlusPlus(\n", " encoder_name='resnet34',\n", " encoder_weights='imagenet',\n", " in_channels=3,\n", " classes=len(CLASS_NAMES),\n", ")\n", "\n", "# ── (Optional) quick training on a toy batch ────────────────────────────────\n", "# Replace this block with your real training loop / checkpoint loading.\n", "optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)\n", "loss_fn = smp.losses.DiceLoss(mode='binary')\n", "\n", "net.train()\n", "for _ in range(3): # 3 fake gradient steps\n", " x = torch.randn(2, 3, IMAGE_SIZE, IMAGE_SIZE)\n", " y = (torch.rand(2, len(CLASS_NAMES), IMAGE_SIZE, IMAGE_SIZE) > 0.5).float()\n", " loss = loss_fn(net(x), y)\n", " optimizer.zero_grad()\n", " loss.backward()\n", " optimizer.step()\n", "\n", "net.eval()\n", "print(f\"Model ready — {sum(p.numel() for p in net.parameters()):,} parameters\")" ] }, { "cell_type": "markdown", "id": "a1ef6a92", "metadata": {}, "source": [ "## 2. Create a DatamintModel Adapter\n", "\n", "Datamint's deployment server calls your `predict_default` method with a list of\n", "`Resource` objects and expects a `list[list[Annotation]]` back — one annotation list per resource.\n", "\n", "You wrap your model by subclassing `DatamintModel` and implementing `predict_default`:\n", "\n", "```\n", "DatamintModel\n", "└── predict_default(model_input: list[Resource], **kwargs) → list[list[Annotation]]\n", "```\n", "\n", "### Key helpers available inside `predict_default`\n", "\n", "| Helper | Returns |\n", "|--------|---------|\n", "| `self.inference_device` | `'cuda'` or `'cpu'` (set by the server) |\n", "| `self.get_pytorch_model()` | The `nn.Module` you passed as `torch_model` |\n", "| `resource.fetch_file_data(auto_convert=True)` | PIL Image / DICOM etc. |\n", "\n", "### Supported annotation types\n", "\n", "Import from `datamint.entities.annotations`:\n", "- `ImageSegmentation` — 2-D binary mask \n", "- `CategoryAnnotation` — class label + optional probability \n", "- `BoxAnnotation` — box from two corner points" ] }, { "cell_type": "code", "execution_count": null, "id": "7dbd4929", "metadata": {}, "outputs": [], "source": [ "import albumentations as A\n", "import cv2\n", "import numpy as np\n", "from albumentations.pytorch import ToTensorV2\n", "\n", "from datamint.entities.annotations import ImageSegmentation\n", "from datamint.mlflow.flavors.model import DatamintModel, ModelSettings\n", "from datamint.mlflow.flavors.task_type import TaskType\n", "\n", "\n", "class SegmentationAdapter(DatamintModel):\n", " \"\"\"Wraps a plain nn.Module for Datamint segmentation inference.\"\"\"\n", "\n", " task_type = TaskType.IMAGE_SEGMENTATION\n", "\n", " def __init__(\n", " self,\n", " torch_model: torch.nn.Module,\n", " class_names: list[str],\n", " image_size: int = 256,\n", " threshold: float = 0.5,\n", " need_gpu: bool = False,\n", " ) -> None:\n", " super().__init__(\n", " torch_model=torch_model,\n", " settings=ModelSettings(need_gpu=need_gpu),\n", " )\n", " self.class_names = class_names\n", " self.image_size = image_size\n", " self.threshold = threshold\n", "\n", " # Preprocessing applied to every image at inference time\n", " self._transform = A.Compose([\n", " A.Resize(image_size, image_size),\n", " A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n", " ToTensorV2(),\n", " ])\n", "\n", " def predict_default(\n", " self,\n", " model_input: list, # list[Resource]\n", " **kwargs,\n", " ) -> list[list[ImageSegmentation]]:\n", " device = self.inference_device\n", " model = self.get_pytorch_model().to(device).eval()\n", "\n", " results = []\n", " for resource in model_input:\n", " # ── Load image ──────────────────────────────────────────────────\n", " img = np.array(resource.fetch_file_data(auto_convert=True, use_cache=True))\n", " if img.ndim == 2: # grayscale → 3-channel\n", " img = np.stack([img, img, img], axis=-1)\n", " elif img.shape[-1] == 4: # RGBA → RGB\n", " img = img[..., :3]\n", " orig_h, orig_w = img.shape[:2]\n", "\n", " # ── Preprocess ──────────────────────────────────────────────────\n", " tensor = self._transform(image=img)['image'].unsqueeze(0).to(device) # (1, 3, H, W)\n", "\n", " # ── Forward pass ────────────────────────────────────────────────\n", " with torch.inference_mode():\n", " logits = model(tensor) # (1, C, H, W)\n", " probs = logits.sigmoid().squeeze(0).cpu().numpy() # (C, H, W) [0, 1]\n", "\n", " # ── Build annotations ───────────────────────────────────────────\n", " annotations = [\n", " ImageSegmentation(\n", " name=self.class_names[i],\n", " segmentation_data=cv2.resize(\n", " (probs[i] > self.threshold).astype(np.uint8),\n", " (orig_w, orig_h),\n", " interpolation=cv2.INTER_NEAREST,\n", " ),\n", " )\n", " for i in range(len(self.class_names))\n", " ]\n", " results.append(annotations)\n", "\n", " return results\n", "\n", "\n", "# Instantiate the adapter with the trained model\n", "adapter = SegmentationAdapter(\n", " torch_model=net,\n", " class_names=CLASS_NAMES,\n", " image_size=IMAGE_SIZE,\n", ")\n", "print(\"Adapter ready.\")" ] }, { "cell_type": "markdown", "id": "4b057353", "metadata": {}, "source": [ "### 2.1 Smoke-test the adapter locally\n", "\n", "Before logging to MLflow it's worth making sure `predict_default` runs without errors\n", "using a fake resource that just wraps a local file." ] }, { "cell_type": "code", "execution_count": null, "id": "3e253b20", "metadata": {}, "outputs": [], "source": [ "import io\n", "\n", "from PIL import Image\n", "\n", "from datamint.entities.resource import LocalResource\n", "\n", "# Create a dummy in-memory PNG as a LocalResource\n", "buf = io.BytesIO()\n", "\n", "img = Image.fromarray(np.random.randint(0, 255, (300, 400, 3), dtype=np.uint8))\n", "img.save(buf, format='PNG')\n", "dummy_resource = LocalResource(raw_data=buf.getvalue(), filename='dummy.png')\n", "\n", "predictions = adapter.predict([dummy_resource])\n", "\n", "print(f\"Received {len(predictions)} result(s)\")\n", "print(f\"Annotations per resource: {[len(p) for p in predictions]}\")\n", "for ann in predictions[0]:\n", " print(f\" {ann.name!r} mask shape={ann.mask.shape} dtype={ann.mask.dtype}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "6376dcce", "metadata": {}, "outputs": [], "source": [ "# upload resource to Datamint (optional, but simulates real usage more closely)\n", "import tempfile\n", "\n", "with tempfile.NamedTemporaryFile(suffix='.png') as tmp:\n", " img.save(tmp.name)\n", " api.resources.upload_resource(tmp.name, tags=['dummy'],\n", " publish_to=proj)" ] }, { "cell_type": "markdown", "id": "687b033a", "metadata": {}, "source": [ "## 3. Log and Register the Model in MLflow\n", "\n", "`datamint.mlflow.flavors.log_model` is the single call that:\n", "\n", "1. Serialises your adapter (including the embedded `nn.Module`) as an MLflow artifact\n", "2. Records the `task_type` so the Datamint server knows how to display predictions\n", "3. Pins the `datamint` and `medimgkit` package versions in `requirements.txt`\n", "4. (Optionally) registers the model in the MLflow Model Registry under a name you choose\n", "\n", "### Why call `datamint.mlflow.set_project` first?\n", "\n", "`set_project` configures the MLflow tracking URI to point to your Datamint server and\n", "associates the run with the correct project, so metrics and artifacts are visible in the\n", "Datamint dashboard." ] }, { "cell_type": "code", "execution_count": null, "id": "f3df96a6", "metadata": {}, "outputs": [], "source": [ "import mlflow\n", "\n", "import datamint.mlflow as datamint_mlflow\n", "from datamint.mlflow.flavors import log_model\n", "\n", "# Point MLflow at the Datamint server and associate runs with the project\n", "datamint_mlflow.set_project(PROJECT_NAME)\n", "mlflow.set_experiment(PROJECT_NAME) # Can be any name; doesn't have to match the project name\n", "\n", "with mlflow.start_run(run_name='external_model_upload') as run:\n", " # Log any metadata you want to track alongside the model\n", " mlflow.log_params({\n", " 'encoder': 'resnet34',\n", " 'image_size': IMAGE_SIZE,\n", " 'class_names': str(CLASS_NAMES),\n", " 'framework': 'segmentation_models_pytorch',\n", " })\n", "\n", " model_info = log_model(\n", " adapter,\n", " task_type=TaskType.IMAGE_SEGMENTATION,\n", " name='segmentation_model', # artifact sub-path inside the run\n", " registered_model_name=MODEL_NAME, # registers in the Model Registry\n", " )\n", "\n", "print(f\"Model URI : {model_info.model_uri}\")\n", "print(f\"Run ID : {run.info.run_id}\")" ] }, { "cell_type": "markdown", "id": "aef7a8a5", "metadata": {}, "source": [ "### 3.1 Assign the `champion` alias\n", "\n", "The deployment API (`api.deploy.start`) resolves models by alias.\n", "Set the `champion` (or any of your choice) alias on the version we just registered." ] }, { "cell_type": "code", "execution_count": null, "id": "45320bce", "metadata": {}, "outputs": [], "source": [ "from mlflow import MlflowClient\n", "\n", "client = MlflowClient()\n", "\n", "# The newly registered version is always the highest version number\n", "versions = client.search_model_versions(f\"name='{MODEL_NAME}'\")\n", "latest_version = max(versions, key=lambda v: int(v.version))\n", "\n", "client.set_registered_model_alias(MODEL_NAME, 'champion', latest_version.version)\n", "print(f\"Alias 'champion' → version {latest_version.version}\")" ] }, { "cell_type": "markdown", "id": "c7ab8c6f", "metadata": {}, "source": [ "## 4. Verify: Load and Run Local Inference\n", "\n", "Before deploying, load the registered model back from MLflow and verify that\n", "`predict` works correctly end-to-end." ] }, { "cell_type": "code", "execution_count": null, "id": "a96affe7", "metadata": {}, "outputs": [], "source": [ "from datamint.mlflow import flavors as datamint_flavor\n", "\n", "model = api.models.get_by_name(MODEL_NAME)\n", "loaded_model = model.get_latest_version(alias='champion').load_model()\n", "\n", "print(f\"Loaded model type : {type(loaded_model).__name__}\")\n", "print(f\"Task type : {loaded_model.task_type}\")\n", "print(f\"Inference device : {loaded_model.inference_device}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "980c4859", "metadata": {}, "outputs": [], "source": [ "# Run inference on a real resource from the project\n", "resources = list(api.resources.get_list(project_name=PROJECT_NAME, limit=3))\n", "\n", "if resources:\n", " predictions = loaded_model.predict(resources[:1])\n", " for ann in predictions[0]:\n", " print(f\" class={ann.name!r} mask={ann.mask.shape} positive_pixels={ann.mask.sum()}\")\n", "else:\n", " # Fall back to the dummy resource from section 2.1\n", " predictions = loaded_model.predict([dummy_resource])\n", " print(\"No resources in project — tested on dummy image.\")\n", " for ann in predictions[0]:\n", " print(f\" class={ann.name!r} mask={ann.mask.shape}\")" ] }, { "cell_type": "markdown", "id": "67b232bd", "metadata": {}, "source": [ "## 5. Deploy to the Datamint Server" ] }, { "cell_type": "code", "execution_count": null, "id": "047ea252", "metadata": {}, "outputs": [], "source": [ "job = api.deploy.start(\n", " model_name=MODEL_NAME,\n", " model_alias='champion',\n", " with_gpu=False,\n", ")\n", "\n", "print(\"Deployment job started\")\n", "print(f\"Job ID : {job.id}\")\n", "print(f\"Status : {job.status}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "0944deb0", "metadata": {}, "outputs": [], "source": [ "import time\n", "\n", "# Poll until the deployment is complete (usually 5–10 minutes)\n", "while True:\n", " job = api.deploy.get_by_id(job.id)\n", " print(f\"Status: {job.status} ({job.progress_percentage:.0f}%)\")\n", " if job.status in ('completed', 'failed', 'cancelled'):\n", " break\n", " time.sleep(15)\n", "\n", "if job.status == 'completed':\n", " print(\"Deployment successful!\")\n", "else:\n", " print(f\"Deployment failed: {job.error_message}\")\n", " print(job.build_logs)" ] }, { "cell_type": "markdown", "id": "3335c541", "metadata": {}, "source": [ "## 6. Remote Inference\n", "\n", "Once deployed, submit resources for server-side inference." ] }, { "cell_type": "code", "execution_count": null, "id": "639935d4", "metadata": {}, "outputs": [], "source": [ "# Pick a resource to run inference on\n", "r = api.resources.get_list(project_name=PROJECT_NAME, limit=1)[0]\n", "\n", "inf_job = api.inference.submit(\n", " model_name=MODEL_NAME,\n", " model_alias='champion',\n", " resource_id=r.id,\n", ")\n", "inf_job.wait() # blocks until inference is complete\n", "\n", "print(f\"Inference complete — {len(inf_job.predictions)} result(s)\")" ] }, { "cell_type": "code", "execution_count": null, "id": "c08a4512", "metadata": {}, "outputs": [], "source": [ "from matplotlib import pyplot as plt\n", "\n", "preds = inf_job.predictions[0]\n", "\n", "fig, axes = plt.subplots(1, len(preds), figsize=(5 * max(len(preds), 1), 5))\n", "if len(preds) == 1:\n", " axes = [axes]\n", "\n", "for ax, ann in zip(axes, preds):\n", " ax.imshow(ann.mask, cmap='gray')\n", " ax.set_title(ann.name)\n", " ax.axis('off')\n", "\n", "plt.suptitle(f\"Remote inference on: {r.filename}\")\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "8dfa590b", "metadata": {}, "source": [ "## 7. Updating the Model (New Version)\n", "\n", "When you retrain or fine-tune your model, log a new version and move the `champion` alias:\n", "\n", "```python\n", "# Re-train / fine-tune your net ...\n", "\n", "adapter_v2 = SegmentationAdapter(torch_model=net_v2, class_names=CLASS_NAMES)\n", "\n", "datamint_mlflow.set_project(PROJECT_NAME)\n", "mlflow.set_experiment(PROJECT_NAME)\n", "\n", "with mlflow.start_run(run_name='external_model_v2'):\n", " mlflow.log_param('version', 2)\n", " model_info = log_model(\n", " adapter_v2,\n", " task_type=TaskType.IMAGE_SEGMENTATION,\n", " name='segmentation_model',\n", " registered_model_name=MODEL_NAME,\n", " )\n", "\n", "# Advance the alias so deployments pick up the new version automatically\n", "versions = client.search_model_versions(f\"name='{MODEL_NAME}'\")\n", "latest_version = max(versions, key=lambda v: int(v.version))\n", "client.set_registered_model_alias(MODEL_NAME, 'champion', latest_version.version)\n", "\n", "# Re-deploy\n", "api.deploy.start(model_name=MODEL_NAME, model_alias='champion', with_gpu=False)\n", "```" ] }, { "cell_type": "markdown", "id": "7d8df686", "metadata": {}, "source": [ "## 8. Custom Prediction Modes (Advanced)\n", "\n", "For large images you may want slice-by-slice or frame-by-frame inference.\n", "Add extra prediction hooks using `@prediction_mode`:\n", "\n", "```python\n", "from datamint.mlflow.flavors.prediction_router import prediction_mode\n", "from datamint.mlflow.flavors.prediction_modes import PredictionMode\n", "\n", "class MyAdapter(DatamintModel):\n", " task_type = TaskType.IMAGE_SEGMENTATION\n", "\n", " def predict_default(self, model_input, **kwargs):\n", " ... # full-image inference\n", "\n", " @prediction_mode(PredictionMode.SLICE)\n", " def predict_slice(self, model_input, slice_idx, **kwargs):\n", " ... # called when mode='slice' is passed as a parameter\n", "```\n", "\n", "Supported modes are defined in `datamint.mlflow.flavors.prediction_modes.PredictionMode`." ] }, { "cell_type": "markdown", "id": "d4c1fb67", "metadata": {}, "source": [ "## Summary\n", "\n", "To deploy any externally trained model with Datamint:\n", "\n", "| Step | Code |\n", "|------|------|\n", "| 1. Wrap your model | Subclass `DatamintModel`, implement `predict_default` |\n", "| 2. Configure MLflow | `datamint.mlflow.set_project(project_name)` |\n", "| 3. Log & register | `log_model(adapter, task_type=..., registered_model_name=...)` |\n", "| 4. Set alias | `client.set_registered_model_alias(name, 'latest', version)` |\n", "| 5. Deploy | `api.deploy.start(model_name=..., model_alias='latest')` |\n", "| 6. Infer | `api.inference.submit(model_name=..., resource_id=...)` |\n", "\n", "The only part unique to your model is the adapter class — everything else is identical\n", "regardless of architecture, framework, or task type." ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.13" } }, "nbformat": 4, "nbformat_minor": 5 }