{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64fddbde",
   "metadata": {},
   "source": [
    "# Fracture Classification with Datamint and FracAtlas Dataset\n",
    "\n",
    "This notebook demonstrates how to build an end-to-end binary classification pipeline using **Datamint** and the **FracAtlas** dataset.\n",
    "\n",
    "## Overview\n",
    "\n",
    "You will learn how to:\n",
    "- **Set up a Datamint project** for managing medical imaging data\n",
    "- **Upload the FracAtlas dataset** to Datamint\n",
    "- **Train a model** using the Datamint Trainer API — one line of code handles dataset loading, model creation, training, and experiment tracking\n",
    "- **Deploy the model** for inference using Datamint's model serving\n",
    "\n",
    "## Comparison with the Manual Workflow\n",
    "\n",
    "The conventional way requires you to:\n",
    "1. Define transforms, a Dataset subclass, and DataLoaders (~80 lines)\n",
    "2. Define a full LightningModule with loss, metrics, and optimizer (~80 lines)\n",
    "3. Configure MLflow logger, callbacks, and Trainer (~30 lines)\n",
    "\n",
    "With the Trainer API, **all of that is replaced by ~3 lines**:\n",
    "\n",
    "```python\n",
    "from datamint.lightning import EfficientNetV2Trainer\n",
    "\n",
    "trainer = EfficientNetV2Trainer(project='FracAtlas')\n",
    "results = trainer.fit()\n",
    "```\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "1. [Setup: Create Project and Upload Dataset](#1-setup-create-project-and-upload-dataset)\n",
    "2. [Training with EfficientNetV2Trainer](#2-training-with-efficientnetv2trainer)\n",
    "3. [Local Inference](#3-local-inference)\n",
    "4. [Custom Model Architecture](#4-custom-model-architecture)\n",
    "5. [Deployment](#5-deployment)\n",
    "\n",
    "## Required Dependencies\n",
    "\n",
    "```bash\n",
    "pip install datamint\n",
    "```\n",
    "\n",
    "## Dataset Overview\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "734f7a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datamint import Api\n",
    "\n",
    "PROJECT_NAME = \"ClassificationTest\"\n",
    "api = Api()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4450f467",
   "metadata": {},
   "source": [
    "## 1. Setup: Create Project and Load Dataset\n",
    "\n",
    "`datamint.examples.fracatlas_dataset` downloads the [FracAtlas dataset](https://doi.org/10.6084/m9.figshare.22363012) (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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "543295f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datamint.examples import fracatlas_dataset\n",
    "\n",
    "proj = fracatlas_dataset.create(project_name=PROJECT_NAME, api=api)\n",
    "proj"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "950ae0f4",
   "metadata": {},
   "source": [
    "### 1.1 Create Train/Validation/Test Splits\n",
    "\n",
    "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).\n",
    "\n",
    "| Split | Percentage | Purpose |\n",
    "|-------|------------|---------|\n",
    "| Train | 80% | Model training |\n",
    "| Validation | 10% | Hyperparameter tuning, early stopping |\n",
    "| Test | 10% | Final model evaluation |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bb794bc8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:datamint.dataset.factory:Detected resource type 'image'; using ImageDataset.\n",
      "INFO:datamint.dataset.base:Allowing external image label '('has_fracture', 'no')' not in project specs.\n",
      "INFO:datamint.dataset.base:Allowing external image label '('has_fracture', 'yes')' not in project specs.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train : 3268 resources\n",
      "val   :  408 resources\n",
      "test  :  409 resources\n"
     ]
    }
   ],
   "source": [
    "from datamint.dataset import build_dataset\n",
    "\n",
    "dataset = build_dataset(\n",
    "    PROJECT_NAME,\n",
    "    include_unannotated=False,\n",
    "    image_categories_merge_strategy='mode',\n",
    "    allow_external_annotations=True,\n",
    ")\n",
    "\n",
    "parts = dataset.split(train=0.8, val=0.1, test=0.1, seed=123)\n",
    "parts.save(force=True)  \n",
    "\n",
    "for name, ds in parts.items():\n",
    "    print(f'{name:6s}: {len(ds):4d} resources')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f18e7a8",
   "metadata": {},
   "source": [
    "## 2. Training with EfficientNetV2Trainer\n",
    "\n",
    "This is where the **Trainer API** shines. Instead of manually defining:\n",
    "- A `Dataset` subclass and DataLoaders\n",
    "- A `LightningModule` with loss, metrics, and optimizer\n",
    "- Callbacks, loggers, and a Lightning `Trainer`\n",
    "- A deployment adapter\n",
    "\n",
    "...you simply create an `EfficientNetV2Trainer` and call `fit()`.\n",
    "\n",
    "The trainer automatically:\n",
    "1. Builds an `ImageDataset` and reads the saved project-scoped splits\n",
    "2. Creates image augmentation pipelines (resize, flip, color jitter, normalize)\n",
    "3. Instantiates **EfficientNetV2-S** pretrained on ImageNet at 384×384\n",
    "4. Sets up **CrossEntropy loss**, **Accuracy** and **F1** metrics\n",
    "5. Configures MLflow logging, checkpointing, and early stopping\n",
    "6. Trains, evaluates on the test split, and **registers the best checkpoint** in the MLflow Model Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4a5dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from datamint.lightning import EfficientNetV2Trainer\n",
    "\n",
    "trainer = EfficientNetV2Trainer(\n",
    "    project=PROJECT_NAME,\n",
    "    image_size=384,\n",
    "    batch_size=8,\n",
    "    max_epochs=3,\n",
    "    accelerator='auto',\n",
    "    pretrained=False,\n",
    "    # model_name='MyModelName'  # defaults to project name\n",
    ")\n",
    "\n",
    "results = trainer.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52547e56",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Test results:\")\n",
    "for metric_dict in results['test_results']:\n",
    "    for k, v in metric_dict.items():\n",
    "        print(f\"  {k}: {v:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "516bb580",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Open the Datamint project dashboard to view experiments and metrics\n",
    "proj.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35c8e762",
   "metadata": {},
   "source": [
    "## Inspect the Registered Model\n",
    "\n",
    "Training automatically registered the model in Datamint's model registry (backed by MLflow).\n",
    "Use `api.models` to inspect it without touching MLflow's client directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af02e20c",
   "metadata": {},
   "outputs": [],
   "source": [
    "registered_model = api.models.get_by_name(PROJECT_NAME)\n",
    "latest_version = registered_model.get_latest_version()\n",
    "\n",
    "print(f\"Version: {latest_version.version}\")\n",
    "print(f\"Task type: {latest_version.get_task_type()}\")\n",
    "print(f\"Supported modes: {latest_version.get_supported_modes()}\")\n",
    "print(\"Metrics:\")\n",
    "for metric_name, value in latest_version.get_metrics().items():\n",
    "    print(f\"  {metric_name}: {value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e977c0",
   "metadata": {},
   "source": [
    "## 3. Local Inference\n",
    "\n",
    "The trainer automatically registers a deployment-ready model in MLflow after training.\n",
    "Load it for local prediction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48591881",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datamint.mlflow import flavors as datamint_flavor\n",
    "\n",
    "# Load the registered model\n",
    "model_loaded = datamint_flavor.load_model(f'models:/{PROJECT_NAME}/latest')\n",
    "\n",
    "# Predict on a resource from the dataset, and log the prediction back to Datamint\n",
    "# as an annotation tagged source='model_deploy' (set log_predictions=True)\n",
    "r = trainer.dataset[0]['resource']\n",
    "prediction = model_loaded.predict([r], params={'log_predictions': True, 'model_name': PROJECT_NAME})\n",
    "print(prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c53d87d",
   "metadata": {},
   "source": [
    "> [!TIP]\n",
    "> You can also serve the model locally via REST API:\n",
    "> ```bash\n",
    "> mlflow models serve -m \"models:/{PROJECT_NAME}/latest\" -p 5111 --env-manager virtualenv\n",
    "> ```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d7ddddd",
   "metadata": {},
   "source": [
    "## 4. Custom Model Architecture\n",
    "\n",
    "The Trainer API lets you swap in a custom backbone while keeping everything else automatic.\n",
    "\n",
    "There are two supported patterns:\n",
    "\n",
    "1. **(Recommended)** Pass a `ClassificationModule` subclass **class** to `model=`. Datamint manages loss, metrics, MLflow logging, and deployment compatibility. Pass the class, not an instance.\n",
    "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`.\n",
    "\n",
    "The example below uses **ResNet-18** from torchvision as a custom backbone while keeping the full Datamint trainer workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecc640a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision.models import resnet18\n",
    "\n",
    "from datamint.lightning import ImageClassificationTrainer\n",
    "from datamint.lightning.trainers.lightning_modules import ClassificationModule\n",
    "\n",
    "\n",
    "class FracAtlasClassifier(ClassificationModule):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        # Replace the default timm model with a custom backbone\n",
    "        self.model = resnet18(weights='DEFAULT')\n",
    "        self.model.fc = torch.nn.Linear(\n",
    "            self.model.fc.in_features, self.hparams['num_classes']\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    # OPTIONAL: override test_step for custom behaviour\n",
    "    # def test_step(self, batch, batch_idx):\n",
    "    #     images = batch['image']              # (B, C, H, W)\n",
    "    #     labels = batch['image_categories']   # (B,)\n",
    "    #     # ... custom logic ...\n",
    "    #     return super().test_step(batch, batch_idx)\n",
    "\n",
    "\n",
    "trainer_custom = ImageClassificationTrainer(\n",
    "    project=PROJECT_NAME,\n",
    "    model=FracAtlasClassifier,   # pass the class, not an instance\n",
    "    image_size=480,\n",
    "    batch_size=8,\n",
    "    max_epochs=10,\n",
    "    accelerator='auto',\n",
    ")\n",
    "\n",
    "results_custom = trainer_custom.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "058f3633",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Test results (custom model):\")\n",
    "for metric_dict in results_custom['test_results']:\n",
    "    for k, v in metric_dict.items():\n",
    "        print(f\"  {k}: {v:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bbfbc81",
   "metadata": {},
   "source": [
    "### 4.1 Using a Fully Custom LightningModule\n",
    "\n",
    "If you already have your own Lightning module, pass an instance. In this mode you own the training and evaluation logic:\n",
    "\n",
    "```python\n",
    "import lightning as L\n",
    "import torch\n",
    "from torchvision.models import resnet50\n",
    "\n",
    "class MyFullyCustomModule(L.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model = resnet50(weights='DEFAULT')\n",
    "        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 2)\n",
    "        self.loss_fn = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        images = batch['image']\n",
    "        labels = batch['image_categories']\n",
    "        loss = self.loss_fn(self(images), labels)\n",
    "        self.log('train/loss', loss)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        images = batch['image']\n",
    "        labels = batch['image_categories']\n",
    "        loss = self.loss_fn(self(images), labels)\n",
    "        self.log('val/loss', loss)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        images = batch['image']\n",
    "        labels = batch['image_categories']\n",
    "        loss = self.loss_fn(self(images), labels)\n",
    "        self.log('test/loss', loss)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters(), lr=1e-4)\n",
    "\n",
    "trainer_full_custom = ImageClassificationTrainer(\n",
    "    project=PROJECT_NAME,\n",
    "    model=MyFullyCustomModule(),  # pass an instance\n",
    "    image_size=480,\n",
    "    batch_size=8,\n",
    "    max_epochs=10,\n",
    "    accelerator='auto',\n",
    ")\n",
    "results = trainer_full_custom.fit()\n",
    "```\n",
    "\n",
    "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`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f673617",
   "metadata": {},
   "source": [
    "## 5. Deployment\n",
    "\n",
    "The trainer **automatically** created and registered a deployment adapter in MLflow after training.\n",
    "Deploy it directly to the Datamint server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3328ca7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = api.deploy.start(\n",
    "    model_name=PROJECT_NAME,\n",
    "    model_alias=\"latest\",\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": "f38939d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = api.deploy.get_by_id(job.id)\n",
    "\n",
    "print(f\"Job Status: {job.status}\")\n",
    "print(f\"Progress: {job.progress_percentage}%\")\n",
    "\n",
    "if job.error_message:\n",
    "    print(f\"Error: {job.error_message}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc8b237",
   "metadata": {},
   "source": [
    "`Model.is_deployed()` reports the same thing, without tracking job IDs, once the job above reaches `completed`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "addd332c",
   "metadata": {},
   "outputs": [],
   "source": [
    "registered_model.is_deployed()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddfabd37",
   "metadata": {},
   "source": [
    "### 5.1 Remote Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2467a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = api.resources.get_list(project_name=PROJECT_NAME, limit=1)[0]\n",
    "\n",
    "inf_job = api.inference.submit(\n",
    "    model_name=PROJECT_NAME,\n",
    "    model_alias=\"latest\",\n",
    "    resource_id=r.id,\n",
    ")\n",
    "inf_job.wait()\n",
    "\n",
    "print(inf_job.predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ddc03fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "905cea65",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "| Step | Description |\n",
    "|------|-------------|\n",
    "| **Data Management** | Uploaded 4,000+ images to Datamint with tags and annotations |\n",
    "| **Splits** | Created 80/10/10 train/val/test splits with `dataset.split().save()` |\n",
    "| **Model Training** | Fine-tuned EfficientNetV2-S with `EfficientNetV2Trainer` — one line |\n",
    "| **Experiment Tracking** | Metrics and model automatically logged with MLflow |\n",
    "| **Custom Model** | Swapped in ResNet-18 by subclassing `ClassificationModule` |\n",
    "| **Deployment** | Model automatically registered and deployed to Datamint server |\n",
    "\n",
    "### See Also\n",
    "\n",
    "- [Datamint Documentation](https://sonanceai.github.io/datamint-python-api/)\n",
    "- [FracAtlas Paper](https://doi.org/10.1038/s41597-023-02432-4)\n",
    "- [PyTorch Lightning Guide](https://lightning.ai/docs/pytorch/stable/)\n",
    "- [MLflow Model Registry](https://mlflow.org/docs/latest/model-registry.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}