diff --git a/.gitignore b/.gitignore index f2ee2044..64035f71 100644 --- a/.gitignore +++ b/.gitignore @@ -112,4 +112,7 @@ junit/ dev/ # vscode -.vscode/ \ No newline at end of file +.vscode/ + +# claude +CLAUDE.md diff --git a/datamint/lightning/__init__.py b/datamint/lightning/__init__.py index d1df3069..97b6bc65 100644 --- a/datamint/lightning/__init__.py +++ b/datamint/lightning/__init__.py @@ -9,6 +9,7 @@ SemanticSegmentation3DTrainer, SegmentationTrainer, UNetPPTrainer, + DeepLabV3PlusTrainer, ) __all__ = [ @@ -20,4 +21,5 @@ "SemanticSegmentation3DTrainer", "SegmentationTrainer", "UNetPPTrainer", + "DeepLabV3PlusTrainer", ] diff --git a/datamint/lightning/trainers/__init__.py b/datamint/lightning/trainers/__init__.py index e2822441..ae8c1dd7 100644 --- a/datamint/lightning/trainers/__init__.py +++ b/datamint/lightning/trainers/__init__.py @@ -6,6 +6,7 @@ from .seg3d_trainer import SemanticSegmentation3DTrainer from .classification_trainer import ClassificationTrainer, ImageClassificationTrainer from .specialized.unetpp import UNetPPTrainer +from .specialized.deeplabv3plus import DeepLabV3PlusTrainer __all__ = [ "BaseTrainer", @@ -13,6 +14,7 @@ "SemanticSegmentation2DTrainer", "SemanticSegmentation3DTrainer", "UNetPPTrainer", + "DeepLabV3PlusTrainer", "ClassificationTrainer", "ImageClassificationTrainer", ] diff --git a/datamint/lightning/trainers/lightning_modules/__init__.py b/datamint/lightning/trainers/lightning_modules/__init__.py index f000be50..2b382e74 100644 --- a/datamint/lightning/trainers/lightning_modules/__init__.py +++ b/datamint/lightning/trainers/lightning_modules/__init__.py @@ -1,6 +1,6 @@ from .base import DatamintLightningModule from .segmentation_module import SegmentationModule -from .segmentation_modules import SMPSegmentationModule, UNetPPModule +from .segmentation_modules import SMPSegmentationModule, UNetPPModule, DeepLabV3PlusModule from .classification_module import ClassificationModule -__all__ = ["DatamintLightningModule", "SegmentationModule", "SMPSegmentationModule", "UNetPPModule", "ClassificationModule"] +__all__ = ["DatamintLightningModule", "SegmentationModule", "SMPSegmentationModule", "UNetPPModule", "DeepLabV3PlusModule", "ClassificationModule"] diff --git a/datamint/lightning/trainers/lightning_modules/segmentation_modules/__init__.py b/datamint/lightning/trainers/lightning_modules/segmentation_modules/__init__.py index 1da35006..d091c6c4 100644 --- a/datamint/lightning/trainers/lightning_modules/segmentation_modules/__init__.py +++ b/datamint/lightning/trainers/lightning_modules/segmentation_modules/__init__.py @@ -1,4 +1,5 @@ from .smp_module import SMPSegmentationModule from .unetpp import UNetPPModule +from .deeplabv3plus import DeepLabV3PlusModule -__all__ = ["SMPSegmentationModule", "UNetPPModule"] +__all__ = ["SMPSegmentationModule", "UNetPPModule", "DeepLabV3PlusModule"] diff --git a/datamint/lightning/trainers/lightning_modules/segmentation_modules/deeplabv3plus.py b/datamint/lightning/trainers/lightning_modules/segmentation_modules/deeplabv3plus.py new file mode 100644 index 00000000..9f54aca7 --- /dev/null +++ b/datamint/lightning/trainers/lightning_modules/segmentation_modules/deeplabv3plus.py @@ -0,0 +1,72 @@ +"""DeepLab v3+ segmentation module.""" +from __future__ import annotations +from collections.abc import Callable +from typing import Any + +import albumentations as A +from torch import Tensor, nn +from typing_extensions import override + +import segmentation_models_pytorch as smp + +from .smp_module import SMPSegmentationModule + + +class DeepLabV3PlusModule(SMPSegmentationModule): + """Segmentation module using the DeepLab v3+ architecture from ``segmentation_models_pytorch``. + + Args: + in_channels: Number of input image channels. + num_classes: Number of segmentation classes excluding background. + loss_fn: Loss module. + metrics_factories: ``{name: callable}`` where each callable returns a fresh metric. + class_names: Human-readable label for each class. + image_size: ``(height, width)`` used during inference. + lr: Learning rate for AdamW. + encoder_name: SMP encoder backbone (e.g. ``'resnet34'``, ``'efficientnet-b4'``). + encoder_weights: Pre-trained weights to load. ``'imagenet'`` by default. + decoder_atrous_rates: Dilation rates for the ASPP module. + Controls the multi-scale receptive field that is DeepLab v3+'s + core architectural feature. SMP default is ``(12, 24, 36)``. + transform: Albumentations transform applied during inference. + """ + + def __init__( + self, + in_channels: int, + num_classes: int, + loss_fn: nn.Module | None = None, + metrics_factories: dict[str, Callable[[], Any]] = {}, + class_names: list[str] | None = None, + image_size: tuple[int, int] | None = None, + lr: float = 1e-4, + encoder_name: str = 'resnet34', + encoder_weights: str | None = 'imagenet', + decoder_atrous_rates: tuple[int, int, int] = (12, 24, 36), + transform: A.BasicTransform | A.BaseCompose | None = None, + ) -> None: + self.in_channels = in_channels + self.num_classes = num_classes + self.encoder_name = encoder_name + self.encoder_weights = encoder_weights + self.decoder_atrous_rates = decoder_atrous_rates + + super().__init__( + transform=transform, + loss_fn=loss_fn, + metrics_factories=metrics_factories, + class_names=class_names, + lr=lr, + ) + + self.model = smp.DeepLabV3Plus( + encoder_name=self.encoder_name, + encoder_weights=self.encoder_weights, + in_channels=self.in_channels, + classes=self.num_classes, + decoder_atrous_rates=self.decoder_atrous_rates, + ) + + @override + def forward(self, x: Tensor) -> Tensor: + return self.model(x) diff --git a/datamint/lightning/trainers/specialized/deeplabv3plus.py b/datamint/lightning/trainers/specialized/deeplabv3plus.py new file mode 100644 index 00000000..1cafa878 --- /dev/null +++ b/datamint/lightning/trainers/specialized/deeplabv3plus.py @@ -0,0 +1,153 @@ +from collections.abc import Callable +from typing import Any, TYPE_CHECKING + +import lightning as L +from torch import nn +from typing_extensions import override + +from ..lightning_modules import DeepLabV3PlusModule +from ..seg2d_trainer import SemanticSegmentation2DTrainer + +if TYPE_CHECKING: + from albumentations import BaseCompose + from datamint.dataset.base import DatamintBaseDataset + from datamint.entities import Project + from medimgkit import ViewPlane + from datamint.lightning.trainers.lightning_modules.base import DatamintLightningModule + + +class DeepLabV3PlusTrainer(SemanticSegmentation2DTrainer): + """Convenience trainer pre-configured for DeepLab v3+. + + Uses the ASPP-based DeepLab v3+ architecture from + ``segmentation_models_pytorch``. The ``decoder_atrous_rates`` parameter + controls the dilation rates of the Atrous Spatial Pyramid Pooling module, + which is DeepLab v3+'s core multi-scale context mechanism. + + Example:: + + trainer = DeepLabV3PlusTrainer( + project='BUS_Segmentation', + encoder_name='resnet50', + ) + results = trainer.fit() + """ + + def __init__( + self, + dataset: 'DatamintBaseDataset | None' = None, + project: 'str | Project | None' = None, + *, + image_size: int | tuple[int, int] | None = None, + slice_axis: 'ViewPlane | int | None' = None, + model: L.LightningModule | type[L.LightningModule] | None = None, + in_channels: int = 3, + loss_fn: nn.Module | None = None, + batch_size: int = 16, + num_workers: int = 4, + train_transform: 'BaseCompose | None' = None, + eval_transform: 'BaseCompose | None' = None, + split_as_of_timestamp: str | None = None, + max_epochs: int = 1, + early_stopping_patience: int | None = 10, + mlflow_experiment_name: str | None = None, + register_model_name: str | None = None, + auto_deploy_adapter: bool = True, + trainer_kwargs: dict[str, Any] | None = None, + dataset_kwargs: dict[str, Any] | None = None, + # DeepLab v3+ specific: + encoder_name: str = 'resnet34', + decoder_atrous_rates: tuple[int, int, int] = (12, 24, 36), + **kwargs: Any, + ) -> None: + """ + Builds a DeepLab v3+ trainer with sensible defaults for segmentation tasks. + + Args: + dataset: A pre-built :class:`DatamintBaseDataset`. Mutually + exclusive with *project*. + project: Project name or :class:`Project` object used to + auto-build a dataset when *dataset* is ``None``. + model: A user-provided :class:`~lightning.LightningModule`. + When ``None`` the trainer builds a default one via + :meth:`_build_model`. + loss_fn: Custom loss function forwarded to the default model. + Ignored when *model* is provided. + batch_size: Training batch size. + num_workers: DataLoader workers. + train_transform: Albumentations transform for training. When + ``None`` the trainer uses :meth:`_train_transform`. + eval_transform: Albumentations transform for val/test. When + ``None`` the trainer uses :meth:`_eval_transform`. + image_size: Target image size ``(H, W)`` or a single int for + square images. Forwarded to default transforms. + split_as_of_timestamp: Historical timestamp used to resolve + project-scoped dataset splits during training. + max_epochs: Maximum number of training epochs. + early_stopping_patience: Epochs without improvement before + stopping. Set to ``None`` to disable early stopping. + mlflow_experiment_name: MLflow experiment name. Auto-generated + from the project name when ``None``. + register_model_name: Name for MLflow Model Registry. + Auto-generated when ``None``. + auto_deploy_adapter: When ``True``, auto-generate a + :class:`~datamint.mlflow.flavors.model.DatamintModel` + adapter after training. + trainer_kwargs: Extra keyword arguments forwarded to + :class:`lightning.Trainer`. + encoder_name: SMP encoder backbone name (e.g. ``'resnet34'``, + ``'resnet50'``, ``'efficientnet-b4'``). + decoder_atrous_rates: Dilation rates for the ASPP module. + Controls the multi-scale receptive field sizes. Defaults to + ``(12, 24, 36)``. Smaller values capture finer-grained context; + larger values capture coarser context. + """ + super().__init__( + dataset=dataset, + project=project, + model=model, + in_channels=in_channels, + image_size=image_size, + slice_axis=slice_axis, + loss_fn=loss_fn, + batch_size=batch_size, + num_workers=num_workers, + train_transform=train_transform, + eval_transform=eval_transform, + split_as_of_timestamp=split_as_of_timestamp, + max_epochs=max_epochs, + early_stopping_patience=early_stopping_patience, + mlflow_experiment_name=mlflow_experiment_name, + register_model_name=register_model_name, + auto_deploy_adapter=auto_deploy_adapter, + trainer_kwargs=trainer_kwargs, + dataset_kwargs=dataset_kwargs, + **kwargs, + ) + self.encoder_name = encoder_name + self.decoder_atrous_rates = decoder_atrous_rates + + @override + def _build_model( + self, + loss_fn: nn.Module, + metrics: dict[str, Callable], + ) -> 'DatamintLightningModule': + num_classes = len(self.dataset.seglabel_list) + if num_classes == 0: + raise ValueError( + "No segmentation labels found in the dataset. " + "DeepLabV3Plus requires at least one segmentation label to train. " + "Make sure your project has annotated resources with segmentation labels, " + "or check that 'include_unannotated' is not masking all annotated data." + ) + return DeepLabV3PlusModule( + encoder_name=self.encoder_name, + in_channels=self.in_channels, + num_classes=num_classes, + loss_fn=loss_fn, + metrics_factories=metrics, + class_names=list(self.dataset.seglabel_list), + image_size=self.image_size, + decoder_atrous_rates=self.decoder_atrous_rates, + ) diff --git a/notebooks/use_cases/segmentation_2d_trainer_BUSI_tutorial.ipynb b/notebooks/use_cases/segmentation_2d_trainer_BUSI_tutorial.ipynb index c3fd8ae2..4596a38e 100644 --- a/notebooks/use_cases/segmentation_2d_trainer_BUSI_tutorial.ipynb +++ b/notebooks/use_cases/segmentation_2d_trainer_BUSI_tutorial.ipynb @@ -7,7 +7,7 @@ "source": [ "# 2D Segmentation with the Datamint Trainer API\n", "\n", - "This notebook shows how to train a **UNet++ semantic segmentation** model on the **BUSI** (Breast Ultrasound Images) dataset using Datamint's **Trainer API** — a high-level wrapper that handles dataset loading, model creation, training, experiment tracking, and model registration from a small amount of code.\n", + "This notebook shows how to train **DeepLab V3+ and UNet++ semantic segmentation** model on the **BUSI** (Breast Ultrasound Images) dataset using Datamint's **Trainer API** — a high-level wrapper that handles dataset loading, model creation, training, experiment tracking, and model registration from a small amount of code.\n", "\n", "## Comparison with the manual Workflow\n", "\n", @@ -19,6 +19,15 @@ "With the Trainer API, **all of that is replaced by ~3 lines**:\n", "\n", "```python\n", + "from datamint.lightning.trainers import DeepLabV3PlusTrainer\n", + "\n", + "trainer = DeepLabV3PlusTrainer(project='MyProject')\n", + "results = trainer.fit()\n", + "```\n", + "\n", + "or\n", + "\n", + "```python\n", "from datamint.lightning.trainers import UNetPPTrainer\n", "\n", "trainer = UNetPPTrainer(project='MyProject')\n", @@ -28,7 +37,7 @@ "## What You'll Learn\n", "\n", "1. Upload data to Datamint\n", - "2. Train with `UNetPPTrainer` using **zero configuration**\n", + "2. Train with `DeepLabV3PlusTrainer` and `UNetPPTrainer` using **zero configuration**\n", "3. Train with `SemanticSegmentation2DTrainer` while swapping in a custom external model\n", "4. Visualise predictions\n", "5. Register and deploy the trained model\n", @@ -49,7 +58,7 @@ "source": [ "from datamint import Api\n", "\n", - "PROJECT_NAME = \"UNetPP_Segmentation_Tutorial\"\n", + "PROJECT_NAME = \"Segmentation_Tutorial\"\n", "api = Api()" ] }, @@ -65,143 +74,14 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "af13262e", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - " \n", - "
\n", - "
Entity
\n", - "
\n", - "

Project

\n", - "
\n", - "
\n", - "\n", - " \n", - "
\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
Name\n", - " UNetPP_Segmentation_Tutorial\n", - "
Created At\n", - " 2025-12-23T14:33:11.807Z\n", - "
Created By\n", - " datamint-dev@mail.com\n", - "
Archived\n", - " False\n", - "
Resource Count\n", - " 780\n", - "
Annotated Resource Count\n", - " 0\n", - "
Description\n", - " Tutorial project for UNet++ segmentation on BTCV dataset\n", - "
\n", - "
\n", - "\n", - "
" - ], - "text/plain": [ - "Project(id='6c770f6c-6483-4af3-924c-16b122bfc6d5', name='UNetPP_Segmentation_Tutorial', created_at='2025-12-23T14:33:11.807Z', created_by='datamint-dev@mail.com', dataset_id='d8c94f56-0897-428b-8d5a-7b846233522a', worklist_id='9ad15be0-4c6d-4537-ab8b-6464fab5eae4', archived=False, resource_count=780, annotated_resource_count=0, description='Tutorial project for UNet++ segmentation on BTCV dataset', viewable_ai_segs=None, editable_ai_segs=None, closed_resources_count=0, resources_to_annotate_count=0, most_recent_experiment=None, annotators=[{'email': 'datamint-dev@mail.com', 'roles': ['PROJECT_OWNER'], 'status': 'active'}])" - ] - }, - "execution_count": 2, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "proj = api.projects.create(\n", " name=PROJECT_NAME,\n", - " description=\"Tutorial project for UNet++ segmentation on BUSI dataset\",\n", + " description=\"Tutorial project for segmentation on BUSI dataset\",\n", " exists_ok=True # Just return the existing project if it already exists\n", ")\n", "proj" @@ -391,7 +271,7 @@ "test_resources = all_resources[n_train + n_val:]\n", "\n", "api.projects.assign_splits(proj, train_resources, split_name='train')\n", - "api.projects.assign_splits(proj, val_resources, split_name='val')\n", + "api.projects.assign_splits(proj, val_resources, split_name='val')### 3.1 UNet++\n", "api.projects.assign_splits(proj, test_resources, split_name='test')\n", "\n", "print(f\"Train: {len(train_resources)}, Val: {len(val_resources)}, Test: {len(test_resources)}\")" @@ -411,17 +291,60 @@ "- Callbacks, loggers, and a Lightning Trainer\n", "- A deployment adapter\n", "\n", - "...you simply create a `UNetPPTrainer` (or `SemanticSegmentation2DTrainer`) and call `fit()`.\n", + "...you simply create a `UNetPPTrainer`, `DeepLabV3PlusTrainer`, or even a `SemanticSegmentation2DTrainer` and call `fit()`.\n", + "\n", "\n", "If your project contains 3D volumes instead of 2D images, the 2D trainer automatically slices them into 2D samples before training.\n", "\n", "The trainer automatically:\n", "1. Builds an `ImageDataset` for 2D projects, or a sliced 2D view of a volume dataset when the project is volumetric\n", - "2. Creates augmentation pipelines (with medical-image-specific augmentations for `UNetPPTrainer`)\n", + "2. Creates augmentation pipelines (with medical-image-specific augmentations for `UNetPPTrainer` (`DeepLabV3PlusTrainer`))\n", "3. Instantiates a default segmentation model when `model=None`\n", "4. Sets up **BCE + Dice loss**, **IoU** and **Dice** metrics\n", "5. Configures MLflow logging, checkpointing, and early stopping\n", - "6. Trains and evaluates on the test split" + "6. Trains and evaluates on the test split.\n", + "\n", + "Here you can choose between two options, DeepLabV3 or UNet++:" + ] + }, + { + "cell_type": "markdown", + "id": "4c06d455", + "metadata": {}, + "source": [ + "### 3.1 DeepLabV3\n", + "\n", + "`DeepLabV3PlusTrainer`\n", + "\n", + "The key architecture-specific parameter is `decoder_atrous_rates`, which controls the dilation rates of the **ASPP** (Atrous Spatial Pyramid Pooling) module — DeepLab v3+'s mechanism for capturing multi-scale context. The default `(12, 24, 36)` works well for most medical imaging tasks." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2c596cd9", + "metadata": {}, + "outputs": [], + "source": [ + "from datamint.lightning import DeepLabV3PlusTrainer\n", + "\n", + "\n", + "trainer = DeepLabV3PlusTrainer(\n", + " project=PROJECT_NAME,\n", + " image_size=256,\n", + " batch_size=16,\n", + " max_epochs=12,\n", + " accelerator='auto',\n", + " # register_model_name='MyModelName' # Default is project name\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b9132010", + "metadata": {}, + "source": [ + "### 3.2 UNet++" ] }, { @@ -444,6 +367,14 @@ ")" ] }, + { + "cell_type": "markdown", + "id": "db8b1870", + "metadata": {}, + "source": [ + "Then, we just need to train our model." + ] + }, { "cell_type": "code", "execution_count": null, @@ -491,7 +422,7 @@ "1. (**RECOMMENDED**) Pass a `SegmentationModule` subclass **class** to `model=`. This is the preferred path when you want Datamint to manage computing loss and metrics, and when you want the resulting model to remain Datamint-compatible for inference and deployment. In this case, pass the class, not the instance object.\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 `DeepLabV3Plus` from `segmentation_models_pytorch` as an external architecture while still keeping the Datamint trainer workflow." + "The example below uses `UNet++` from `segmentation_models_pytorch` as an external architecture while still keeping the Datamint trainer workflow." ] }, { @@ -591,7 +522,7 @@ "import mlflow\n", "from mlflow import MlflowClient\n", "\n", - "model_uri = \"models:/UNetPP_Segmentation_Tutorial/latest\"\n", + "model_uri = \"models:/Segmentation_Tutorial/latest\"\n", "\n", "model_info = mlflow.models.get_model_info(model_uri)\n", "client = MlflowClient()\n", @@ -635,11 +566,11 @@ "\n", "r = trainer.dataset[0]['resource']\n", "\n", - "model_loaded = datamint_flavor.load_model('models:/UNetPP_Segmentation_Tutorial/latest')\n", + "model_loaded = datamint_flavor.load_model('models:/Segmentation_Tutorial/latest')\n", "model_loaded.predict([r])\n", "\n", "# Alternatively:\n", - "# model_loaded = mlflow.pyfunc.load_model('models:/UNetPP_Segmentation_Tutorial/latest')\n", + "# model_loaded = mlflow.pyfunc.load_model('models:/Segmentation_Tutorial/latest')\n", "# model_loaded.predict([r])" ] }, @@ -736,7 +667,7 @@ "\n", "The Trainer API is customizable at several layers. Here are the most common overrides:\n", "\n", - "### 5.1 Change the encoder backbone\n", + "### 5.1 Change the encoder backbone (example using uNet++)\n", "\n", "Use any encoder supported by [segmentation_models_pytorch](https://github.com/qubvel-org/segmentation_models.pytorch):\n", "\n", @@ -930,7 +861,7 @@ "outputs": [], "source": [ "from datamint import Api\n", - "PROJECT_NAME = \"UNetPP_Segmentation_Tutorial\"\n", + "PROJECT_NAME = \"Segmentation_Tutorial\"\n", "api = Api()" ] },