diff --git a/docs/notebooks/segmentation_pipeline.ipynb b/docs/notebooks/segmentation_pipeline.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..4983a7fd65fe73c18a5b270e824372d554eb41c9
--- /dev/null
+++ b/docs/notebooks/segmentation_pipeline.ipynb
@@ -0,0 +1,418 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Deep learning volume segmentation\n",
+ "\n",
+ "Authors: Alessia Saccardo (s212246@dtu.dk) & Felipe Delestro (fima@dtu.dk)\n",
+ "\n",
+ "This notebook aims to demonstrate the feasibility of implementing a comprehensive deep learning segmentation pipeline solely leveraging the capabilities offered by the qim3d library. Specifically, it will highlight the utilization of the annotation tool and walk through the process of creating and training a Unet model."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Imports"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import qim3d\n",
+ "import numpy as np \n",
+ "import os "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Load data\n",
+ "Qim3d library contains a set of example volumes which can be easily loaded using `qim3d.examples.{volume_name}`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "vol = qim3d.examples.bone_128x128x128"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "To easily have an insight of how the volume looks like we can interact with it using the `slicer` function from `qim3d`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "qim3d.viz.slicer(vol)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Generate dataset for training\n",
+ "\n",
+ "In order to train the classification model, we need to create a dataset from the volume.\n",
+ "\n",
+ "This means that we'll need a few slices to be used for `training` and at least one for the `test`"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The dataset for training the model is managed by `qim3d.utils.prepare_datasets` and it expects files to follow this structure:\n",
+ "\n",
+ "<pre>\n",
+ "dataset\n",
+ "├── test\n",
+ "│ ├── images\n",
+ "│ │ └── FileA.png\n",
+ "│ └── labels\n",
+ "│ └── FileA.png\n",
+ "└── train\n",
+ " ├── images\n",
+ " │ ├── FileB.png\n",
+ " │ └── FileC.png\n",
+ " └── labels\n",
+ " ├── FileB.png\n",
+ " └── FileC.png\n",
+ "\n",
+ "\n",
+ "</pre>\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Number of slices that will be used\n",
+ "ntraining = 4\n",
+ "ntest = 1"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "In the following cell, we get the slice indices, making sure that we're not using the same indices for training and test."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Create a set with all the indices\n",
+ "all_idxs = set(range(vol.shape[0]))\n",
+ "\n",
+ "# Get indices for training data\n",
+ "training_idxs = list(np.random.choice(list(all_idxs), size=ntraining))\n",
+ "print(f\"Slices for training data...: {training_idxs}\")\n",
+ "\n",
+ "# Get indices for test data\n",
+ "test_idxs = list(np.random.choice(list(all_idxs - set(training_idxs)), size=ntest, replace=False))\n",
+ "print(f\"Slices for test data.......: {test_idxs}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Create folder structure\n",
+ "Here we create the necessary directories"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Base path for the training data\n",
+ "base_path = os.path.expanduser(\"~/dataset\")\n",
+ "\n",
+ "# Create directories\n",
+ "print(\"Creating directories:\")\n",
+ "for folder_split in [\"train\", \"test\"]:\n",
+ " for folder_type in [\"images\", \"labels\"]:\n",
+ " path = os.path.join(base_path, folder_split, folder_type)\n",
+ " os.makedirs(path, exist_ok=True)\n",
+ " print(path)\n",
+ "\n",
+ "# Here we have the option to remove any previous files\n",
+ "clean_files = True\n",
+ "if clean_files:\n",
+ " for root, dirs, files in os.walk(base_path):\n",
+ " for file in files:\n",
+ " file_path = os.path.join(root, file)\n",
+ " os.remove(file_path)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Annotate data\n",
+ "\n",
+ "The following cell will generate an annotation tool for each slice that was requested. \n",
+ "\n",
+ "You should use the tool to drawn a mask over the structures you're willing to detect, and press the button `Update` so that the mask is saved."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "annotation_tools = {}\n",
+ "for idx in training_idxs + test_idxs:\n",
+ " if idx in training_idxs:\n",
+ " subset = \"training\"\n",
+ " elif idx in test_idxs:\n",
+ " subset = \"test\"\n",
+ " annotation_tools[idx] = qim3d.gui.annotation_tool.Interface()\n",
+ " annotation_tools[idx].name_suffix = f\"_{idx}\" \n",
+ " print(f\"Annotation for slice {idx} ({subset})\")\n",
+ " annotation_tools[idx].launch(vol[idx])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Getting masks from the annotation tool\n",
+ "The masks are stored in the annotation tool when the button `Update` is pressed\n",
+ "\n",
+ "Here we extract the masks and save them to disk, followign the standard needed for the DL model"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(\"Saving images and masks to disk\")\n",
+ "for idx in training_idxs + test_idxs:\n",
+ " \n",
+ " if idx in training_idxs:\n",
+ " folder_split = \"train\"\n",
+ "\n",
+ " elif idx in test_idxs:\n",
+ " folder_split = \"test\"\n",
+ "\n",
+ " print (f\"- slice {idx} ({folder_split})\")\n",
+ " mask_dict = annotation_tools[idx].get_result()\n",
+ " mask = list(mask_dict.values())[0]\n",
+ "\n",
+ " # Save image\n",
+ " qim3d.io.save(os.path.join(base_path,folder_split,\"images\",f\"{idx}.png\"), vol[idx], replace=True)\n",
+ " # Save label\n",
+ " qim3d.io.save(os.path.join(base_path,folder_split,\"labels\",f\"{idx}.png\"), mask, replace=True)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Build Unet"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Building and training the Unet model is straightforward using `qim3d`. \n",
+ "\n",
+ "We first need to instantiate the model by defying its size, which can be either *small*, *medium* or *large*. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# defining model\n",
+ "model = qim3d.models.UNet(size = 'small', dropout = 0.25)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Then we need to decide which type of augumentation to apply to the data. \n",
+ "\n",
+ "The `qim3d.utils.Augmentation` allows to specify how the images should be reshaped to the appropriate size and the level of transformation to apply respectively to train, test and validation sets. \n",
+ "\n",
+ "The resize must be choosen between [*crop*, *reshape*, *padding*] and the level of transformation must be chosse between [*None*, *light*, *moderate*, *heavy*]. The user can also specify the mean and standard deviation values for normalizing pixel intensities."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# defining augmentation\n",
+ "aug = qim3d.utils.Augmentation(resize = 'crop', transform_train = 'light')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Then the datasets and dataloaders are instantiated "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# datasets and dataloaders\n",
+ "train_set, val_set, test_set = qim3d.utils.prepare_datasets(path = base_path,\n",
+ " val_fraction = 0.5,\n",
+ " model = model,\n",
+ " augmentation = aug)\n",
+ "\n",
+ "\n",
+ "train_loader, val_loader, test_loader = qim3d.utils.prepare_dataloaders(train_set, \n",
+ " val_set,\n",
+ " test_set,\n",
+ " batch_size = 1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "The hyperparameters are defined using the function `qim3d.models.Hyperparameters` and the model can be easily trained by running the function `qim3d.utils.train_model` which returns also a plot of the losses at the end of the training if the option is selected by the user "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Train model"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# model hyperparameters\n",
+ "hyperparameters = qim3d.models.Hyperparameters(model, n_epochs=10, \n",
+ " learning_rate = 5e-3, loss_function='DiceCE',\n",
+ " weight_decay=1e-3)\n",
+ "\n",
+ "# training model\n",
+ "qim3d.utils.train_model(model, hyperparameters, train_loader, val_loader, plot=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Check results\n",
+ "\n",
+ "To compute the inference step it is just needed to run `qim3d.utils.inference`.\n",
+ "\n",
+ "The results can be visualize with the function `qim3d.viz.grid_pred` that shows the predicted segmentation along with a comparison between the ground truth."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "in_targ_preds_test = qim3d.utils.inference(test_set, model)\n",
+ "qim3d.viz.grid_pred(in_targ_preds_test,alpha=1)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Compute inference on entire volume\n",
+ "\n",
+ "Given that the input is a volume, the goal is to perform inference on the entire volume rather than individual slices.\n",
+ "\n",
+ "By using the function `qim3d.utils.volume_inference` it is possible to obtain the segmentation volume output"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "inference_vol = qim3d.utils.models.volume_inference(vol, model)\n",
+ "qim3d.viz.slicer(inference_vol)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "We can also visualize the created mask together with the original volume"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "vol_masked = qim3d.viz.vol_masked(vol, inference_vol, viz_delta=128)\n",
+ "qim3d.viz.slicer(vol_masked, cmap=\"PiYG\")"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "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.11.5"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/qim3d/gui/annotation_tool.py b/qim3d/gui/annotation_tool.py
index 9580739914514e9eb47bbc02da5a0f27c20816d7..b39e06a2c583d0de6712108afece78c6a5db4258 100644
--- a/qim3d/gui/annotation_tool.py
+++ b/qim3d/gui/annotation_tool.py
@@ -48,6 +48,7 @@ class Interface:
self.interface = None
self.username = getpass.getuser()
self.temp_dir = os.path.join(tempfile.gettempdir(), f"qim-{self.username}")
+ self.name_suffix = None
# CSS path
current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -77,7 +78,7 @@ class Interface:
# Get the temporary files from gradio
temp_path_list = []
for filename in os.listdir(self.temp_dir):
- if "mask" in str(filename):
+ if "mask" and self.name_suffix in str(filename):
# Get the list of the temporary files
temp_path_list.append(os.path.join(self.temp_dir, filename))
@@ -116,7 +117,15 @@ class Interface:
with gr.Column(scale=6):
img_editor = gr.ImageEditor(
# ! Temporary fix for drawing at wrong location https://github.com/gradio-app/gradio/pull/7959
- value={"background": img, "layers": [Image.new("RGBA", img.shape, (0, 0, 0, 0))], "composite": None} if img is not None else None,
+ value=(
+ {
+ "background": img,
+ "layers": [Image.new("RGBA", img.shape, (0, 0, 0, 0))],
+ "composite": None,
+ }
+ if img is not None
+ else None
+ ),
type="numpy",
image_mode="RGB",
brush=brush,
@@ -147,12 +156,14 @@ class Interface:
)
temp_dir = gr.Textbox(value=self.temp_dir, visible=False)
+ name_suffix = gr.Textbox(value=self.name_suffix, visible=False)
+
session = gr.State([])
inputs = [img_editor]
operations = Operations()
# fmt: off
btn_update.click(
- fn=operations.start_session, inputs=[img_editor,temp_dir] , outputs=session).then(
+ fn=operations.start_session, inputs=[img_editor,temp_dir, name_suffix] , outputs=session).then(
fn=operations.preview, inputs=session, outputs=overlay_img).then(
fn=self.set_visible, inputs=None, outputs=overlay_img).then(
fn=operations.separate_masks, inputs=session, outputs=[session, masks_download]).then(
@@ -168,13 +179,18 @@ class Operations:
session = Session()
session.img_editor = args[0]
session.temp_dir = args[1]
+ session.mask_names = {
+ 0: f"red{args[2]}",
+ 1: f"green{args[2]}",
+ 2: f"blue{args[2]}",
+ }
# Clean up old files
try:
files = os.listdir(session.temp_dir)
for filename in files:
# Check if "mask" is in the filename
- if "mask" in filename:
+ if "mask" and args[2] in filename:
file_path = os.path.join(session.temp_dir, filename)
os.remove(file_path)
diff --git a/qim3d/utils/models.py b/qim3d/utils/models.py
index 19a2a844fe225932e03b041c922328129aa771e8..82cfc78918eda7d9e7d84a9253b1d6b36aaa8c61 100644
--- a/qim3d/utils/models.py
+++ b/qim3d/utils/models.py
@@ -1,5 +1,6 @@
""" Tools performed with models."""
import torch
+import numpy as np
import matplotlib.pyplot as plt
from torchinfo import summary
@@ -227,4 +228,38 @@ def inference(data,model):
targets = targets.unsqueeze(0)
preds = preds.unsqueeze(0)
- return inputs,targets,preds
\ No newline at end of file
+ return inputs,targets,preds
+
+
+def volume_inference(volume, model, threshold=0.5):
+ '''
+ Compute on the entire volume
+ Args:
+ volume (numpy.ndarray): A 3D numpy array representing the input volume.
+ model (torch.nn.Module): The trained network model used for inference.
+ threshold (float): The threshold value used to binarize the model predictions.
+ Returns:
+ numpy.ndarray: A 3D numpy array representing the model predictions for each slice of the input volume.
+ Raises:
+ ValueError: If the input volume is not a 3D numpy array.
+ '''
+ if len(volume.shape) != 3:
+ raise ValueError("Input volume must be a 3D numpy array")
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ model.to(device)
+ model.eval()
+
+ inference_vol = np.zeros_like(volume)
+
+ for idx in np.arange(len(volume)):
+ input_with_channel = np.expand_dims(volume[idx], axis=0)
+ input_tensor = torch.tensor(input_with_channel, dtype=torch.float32).to(device)
+ input_tensor = input_tensor.unsqueeze(0)
+ output = model(input_tensor) > threshold
+ output = output.cpu() if device == 'cuda' else output
+ output_detached = output.detach()
+ output_numpy = output_detached.numpy()[0, 0, :, :]
+ inference_vol[idx] = output_numpy
+
+ return inference_vol
\ No newline at end of file
diff --git a/qim3d/viz/__init__.py b/qim3d/viz/__init__.py
index 1a9c76da2709f8f39ed417a481c62427235db994..c11df18022667389276dfafaf077b5579f388801 100644
--- a/qim3d/viz/__init__.py
+++ b/qim3d/viz/__init__.py
@@ -1,5 +1,5 @@
from .visualizations import plot_metrics
-from .img import grid_pred, grid_overview, slices, slicer, orthogonal
+from .img import grid_pred, grid_overview, slices, slicer, orthogonal, vol_masked
from .k3d import vol
from .structure_tensor import vectors
from .local_thickness_ import local_thickness
diff --git a/qim3d/viz/img.py b/qim3d/viz/img.py
index 4acefda157683a86b6f59447d2e2cfb3d106c159..d06697be9e0772135159026d27d0fa5c1e6dc23f 100644
--- a/qim3d/viz/img.py
+++ b/qim3d/viz/img.py
@@ -15,7 +15,6 @@ from matplotlib.colors import LinearSegmentedColormap
from qim3d.io.logger import log
-
def grid_overview(
data, num_images=7, cmap_im="gray", cmap_segm="viridis", alpha=0.5, show=False
):
@@ -335,7 +334,9 @@ def slices(
slice_idx = i * max_cols + j
try:
slice_img = vol.take(slice_idxs[slice_idx], axis=axis)
- ax.imshow(slice_img, cmap=cmap, interpolation=interpolation, **imshow_kwargs)
+ ax.imshow(
+ slice_img, cmap=cmap, interpolation=interpolation, **imshow_kwargs
+ )
if show_position:
ax.text(
@@ -520,4 +521,33 @@ def orthogonal(
y_slicer.children[0].description = "Y"
x_slicer.children[0].description = "X"
- return widgets.HBox([z_slicer, y_slicer, x_slicer])
\ No newline at end of file
+ return widgets.HBox([z_slicer, y_slicer, x_slicer])
+
+
+def vol_masked(vol, vol_mask, viz_delta=128):
+ """
+ Applies masking to a volume based on a binary volume mask.
+
+ This function takes a volume array `vol` and a corresponding binary volume mask `vol_mask`.
+ It computes the masked volume where pixels outside the mask are set to the background value,
+ and pixels inside the mask are set to foreground.
+
+
+ Args:
+ vol (ndarray): The input volume as a NumPy array.
+ vol_mask (ndarray): The binary mask volume as a NumPy array with the same shape as `vol`.
+ viz_delta (int, optional): Value added to the volume before applying the mask to visualize masked regions.
+ Defaults to 128.
+
+ Returns:
+ ndarray: The masked volume with the same shape as `vol`, where pixels outside the mask are set
+ to the background value (negative).
+
+
+ """
+
+ background = (vol.astype("float") + viz_delta) * (1 - vol_mask) * -1
+ foreground = (vol.astype("float") + viz_delta) * vol_mask
+ vol_masked = background + foreground
+
+ return vol_masked