{ "cells": [ { "cell_type": "markdown", "id": "a1b2c3d4", "metadata": {}, "source": [ "# 3D Segmentation with UNETR++ and the Datamint Trainer API\n", "\n", "This notebook shows how to train a **UNETR++** volumetric segmentation model on the **Synapse Multi-Organ CT** (https://www.kaggle.com/datasets/dogcdt/synapse) dataset using Datamint's **Trainer API**.\n", "\n", "## What is UNETR++?\n", "\n", "UNETR++ (Shaker et al., 2024) is a **true 3D** segmentation model, it takes full volumetric inputs `(B, C, D, H, W)` and outputs 3D segmentation maps, unlike slice-based approaches that process 2D cuts independently.\n", "\n", "The architecture:\n", "- **Encoder**: 4-stage hierarchical transformer with **Efficient Paired Attention (EPA)**, decomposes attention into a *channel branch* and a *spatial branch* to keep complexity sub-quadratic\n", "- **Decoder**: 4 transposed-conv blocks with element-wise skip connections; the final block uses a ResBlock (not a transformer) since full-resolution attention is too expensive\n", "\n", "> Shaker et al., \"UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation\", IEEE TMI 2024\n", "\n", "## What You'll Learn\n", "\n", "1. Download and preprocess the Synapse Multi-Organ CT dataset\n", "2. Upload 3D volumes and segmentation masks to Datamint\n", "3. Train with `UNETRPPTrainer` using a demo-friendly config\n", "4. Visualise 3D predictions on axial slices\n", "\n", "## Required Dependencies\n", "\n", "```bash\n", "pip install datamint gdown nibabel h5py\n", "```\n", "\n", "## ⚠️ Dataset Access\n", "\n", "The Synapse dataset can be found at Kaggle." ] }, { "cell_type": "code", "execution_count": 1, "id": "b2c3d4e5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Note: you may need to restart the kernel to use updated packages.\n" ] } ], "source": [ "%pip install -q datamint h5py nibabel" ] }, { "cell_type": "code", "execution_count": null, "id": "c3d4e5f6", "metadata": {}, "outputs": [], "source": [ "from datamint import Api\n", "\n", "PROJECT_NAME = \"UNETRPP_Synapse_Tutorial\"\n", "api = Api()" ] }, { "cell_type": "markdown", "id": "d4e5f6a7", "metadata": {}, "source": [ "## 1. Setup: Create Project and Load Dataset\n", "\n", "**Label map** (8 abdominal organs):\n", "\n", "| Value | Organ | Value | Organ |\n", "|-------|-------|-------|-------|\n", "| 1 | aorta | 5 | right_kidney |\n", "| 2 | gallbladder | 6 | liver |\n", "| 3 | spleen | 7 | stomach |\n", "| 4 | left_kidney | 8 | pancreas |\n", "\n", "`datamint.examples.synapse_dataset` downloads the [Synapse Multi-Organ CT dataset](https://www.kaggle.com/datasets/dogcdt/synapse), converts each HDF5 volume to NIfTI (`.nii.gz`, downsampled 2x in-plane), and uploads the CT volumes and organ segmentation masks to a Datamint project. The call is idempotent, if the project already exists, it's returned as-is without re-uploading." ] }, { "cell_type": "code", "execution_count": null, "id": "e5f6a7b8", "metadata": {}, "outputs": [], "source": [ "from datamint.examples import synapse_dataset\n", "\n", "proj = synapse_dataset.create(project_name=PROJECT_NAME, api=api)\n", "proj" ] }, { "cell_type": "markdown", "id": "b4c5d6e7", "metadata": {}, "source": [ "### 1.1 Assign Train / Val / Test Splits" ] }, { "cell_type": "code", "execution_count": 6, "id": "c5d6e7f8", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Train: 8 Val: 1 Test: 3\n" ] } ], "source": [ "import random\n", "\n", "all_resources = list(api.resources.get_list(project_name=PROJECT_NAME))\n", "all_resources.sort(key=lambda r: r.filename)\n", "\n", "random.seed(42)\n", "random.shuffle(all_resources)\n", "\n", "n = len(all_resources)\n", "n_train = int(0.70 * n)\n", "n_val = int(0.15 * n)\n", "\n", "api.projects.assign_splits(all_resources[:n_train], split_name='train', project=proj)\n", "api.projects.assign_splits(all_resources[n_train:n_train + n_val], split_name='val', project=proj)\n", "api.projects.assign_splits(all_resources[n_train + n_val:], split_name='test', project=proj)\n", "\n", "print(f\"Train: {n_train} Val: {n_val} Test: {n - n_train - n_val}\")" ] }, { "cell_type": "markdown", "id": "d6e7f8a9", "metadata": {}, "source": [ "## 2. Train UNETR++\n", "\n", "We use a **demo config** sized for ≥ 4 GB VRAM." ] }, { "cell_type": "code", "execution_count": 4, "id": "e7f8a9b0", "metadata": {}, "outputs": [], "source": [ "\n", "from datamint.lightning import UNETRPPTrainer\n", "\n", "trainer = UNETRPPTrainer(\n", " project=PROJECT_NAME,\n", " patch_crop_size=(64, 64, 64),\n", " feature_size=16,\n", " num_heads=4,\n", " depths=[1, 1, 1, 1],\n", " in_channels=1,\n", " batch_size=1,\n", " max_epochs=3,\n", " sw_overlap=0.5,\n", " early_stopping_patience=1,\n", " accelerator='gpu',\n", ")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f8a9b0c1", "metadata": {}, "outputs": [], "source": [ "results = trainer.fit()" ] }, { "cell_type": "markdown", "id": "fd50d89b", "metadata": {}, "source": [ "