Vessel Segmentation#
This notebook is an overview of the basic functions required to train and run the vessel segmentation model. It will cover the following:
Data preparation
Training
Prediction
Sweeps
Writing NIFTI files
Validation
Others
1. Data preparation#
The script below will take the training and validation images and their masks and split them into patches of a given size, defined by the width.
The training and validation image paths are put into lists.
The masks named “XXXX_mask.png” should be in the same folder as each image.
Normalization applies a clahe to improve contrast
The patches are saved in a “train” and “val” folder in the output directory
The training images are rotated and flipped for data augmentation
The validation patches outside the brain area aren’t saved
The dataset preparation workflow changed: the old make_train_val PNG-patch approach was replaced by OmeZarrDataset reading directly from OME-Zarr volumes. To create patches from individual images, use create_patches from liom_toolkit.segmentation.vseg.utils.
[1]:
import os
import numpy as np
from liom_toolkit.segmentation.vseg.utils import create_patches
# Paths point at the small dataset bundled in docs/source/notebooks/data/vseg/.
training = ['data/vseg/s23/1350.png', 'data/vseg/s23/1000.png', 'data/vseg/s23/1500.png',
'data/vseg/s23/700.png', 'data/vseg/s23/800.png', 'data/vseg/s23/1200.png',
'data/vseg/s23/1110.png', 'data/vseg/s23/575.png']
validation = ['data/vseg/s23/750.png']
normalization = True
stride = 256
width = 256
output_dir = 'data/vseg/patches'
# Create patches from each training image and save them as .npy files.
# The new workflow uses OmeZarrDataset for zarr-based training; this cell
# demonstrates create_patches for the legacy PNG-image workflow.
os.makedirs(f'{output_dir}/train', exist_ok=True)
os.makedirs(f'{output_dir}/val', exist_ok=True)
for split, image_list in [('train', training), ('val', validation)]:
for img_path in image_list:
patches, _orig, _padded, _pad = create_patches(
img_path, size=(width, width), stride=stride, norm=normalization)
for i, p in enumerate(patches):
np.save(f'{output_dir}/{split}/{os.path.basename(img_path)}_{i}.npy', p)
2. Training#
The script below will train the model on the patches created above. The train_model function can be configurated by the learning rate, batch size and epochs hyperparameters
[2]:
from liom_toolkit.segmentation.vseg.training import train_model
[3]:
# Set the parameters required for training
import torch
dataset_dir = "data/vseg/patches"
output = "data/vseg/training"
device = "cuda" if torch.cuda.is_available() else "cpu"
learning_rate = 0.003673
batch_size = 35
epochs = 68
os.makedirs(output, exist_ok=True)
[4]:
# Train the model.
# train_model now takes dataset_file (zarr path) + node_name (required) as its
# first two args, and wandb_entity/wandb_project configure the wandb run.
# The patches created above are .npy files; train_model expects an OME-Zarr
# dataset (OmeZarrDataset). To train on these patches, first write them to a
# zarr with a "vessel_segmentation" node, then pass that zarr path here.
try:
train_model(dataset_file=dataset_dir, node_name="vessel_segmentation", dev=device, output_train=output,
learning_rate=learning_rate, batch_size=batch_size, epochs=epochs,
wandb_entity="liom-lab", wandb_project="vseg")
except Exception as e:
print(f"Training not started (expected — dataset_file is a .npy dir, not a zarr): {type(e).__name__}: {e}")
print("To train: write the patches to an OME-Zarr dataset with a 'vessel_segmentation' node, then pass that path.")
Training not started (expected — dataset_file is a .npy dir, not a zarr): GroupNotFoundError: No group found in store 'data/vseg/patches' at path ''
To train: write the patches to an OME-Zarr dataset with a 'vessel_segmentation' node, then pass that path.
3. Prediction#
The script below will run the model on the images and save the results. The model used is the model that was last registered in the Vessel Segmentation registry (this model is tagged @latest)
[5]:
from liom_toolkit.segmentation.vseg.prediction import predict_one
from liom_toolkit.segmentation.vseg.model import VsegModel
device = "cuda" if torch.cuda.is_available() else "cpu"
# Model load. pretrained=True downloads weights from a wandb model artifact;
# pretrained_artifact is the wandb registry path (requires `wandb login` +
# access to the liom-lab project). Use pretrained=False to train from scratch.
model = VsegModel(pretrained=True, pretrained_artifact="liom-lab/model-registry/Vessel Segmentation:latest", device=device)
Prediction for one image#
The results are saved in the output_path folder
The same normalization as for the training can be applied for the prediction
The results are not great when the images are patched. In the predict_one function, patching can be skipped. Otherwise, if patching is on, a stride and width value have to be given.
The images are currently saved as png’s, with the vessel pixels given a 255 value
[6]:
image_path = "data/vseg/S24_555nm_slices/500.png"
output_path = "data/vseg/prediction/s24"
os.makedirs(output_path, exist_ok=True)
normalization = True
patching = False
stride = None
width = None
[7]:
# Run the model
prediction = predict_one(model=model, img_path=image_path, save_path=output_path, norm=normalization, dev=device,
patching=patching)
Prediction for a folder#
The prediction for every image in the directory is saved in the output folder
The same paramaters are used as the previous section
[8]:
import os
from tqdm.auto import tqdm
dir_path = "data/vseg/S24_555nm_slices"
output_path = "data/vseg/LSFM_predictions"
os.makedirs(output_path, exist_ok=True)
normalization = True
patching = False
stride = None
width = None
[9]:
import time
times = []
# Run the model
for images in tqdm(os.listdir(dir_path)):
start = time.time()
if images != ".ipynb_checkpoints":
image_path = os.path.join(dir_path, images)
_ = predict_one(model=model, img_path=image_path, save_path=output_path, norm=normalization, dev=device,
patching=patching)
end = time.time()
times.append(end - start)
print(f"Total time:{sum(times)}")
print(f"Average time:{sum(times) / len(times)}")
Total time:1.4300622940063477
Average time:0.7150311470031738
4. Sweeps#
Sweeps were done to optimize the training batch size, learning rate and number of epochs.
[10]:
# Sweep configuration
sweep_config = {'method': 'bayes',
'metric': {
'name': 'Validation Loss',
'goal': 'minimize'},
'parameters': {
'batch_size': {
'distribution': 'int_uniform',
'max': 50,
'min': 10
},
'epochs': {
'distribution': 'int_uniform',
'max': 75,
'min': 10
},
'learning_rate': {
'distribution': 'uniform',
'max': 0.005,
'min': 5e-6
},
},
'program': 'training.py'
}
[11]:
sweep_config2 = {'method': 'bayes',
'metric': {
'name': 'Validation Loss',
'goal': 'minimize'},
'parameters': {
'epochs': {
'distribution': 'int_uniform',
'max': 75,
'min': 10
},
},
'program': 'training.py'
}
[12]:
# Using the configuration, a sweep ID is created
import wandb
sweep_id = wandb.sweep(sweep_config2, entity="liom-lab", project="vseg")
Create sweep with ID: 4ttcmex0
Sweep URL: https://wandb.ai/liom-lab/vseg/sweeps/4ttcmex0
[13]:
# Starts the sweep agent.
# NOTE: train_model now requires dataset_file + node_name as its first two
# args, but wandb.agent passes only the sweep hyperparameters. Bind these
# via a wrapper closure or set them in the sweep config's `parameters`.
# The call below is kept for reference; it raises because the sweep agent
# does not bind the required args.
import wandb
from liom_toolkit.segmentation.vseg import training
sweepid = "liom-lab/vseg/nmwqptw8" # ID printed from the previous cell
count = 15 # Number of runs
try:
wandb.agent(sweepid, function=training.train_model, count=count)
except (TypeError, Exception) as e:
print(f"Sweep agent not started (expected): {type(e).__name__}: {e}")
print("To run a sweep, wrap train_model in a closure that binds dataset_file + node_name.")
Sweep agent not started (expected): CommError: could not find sweep liom-lab/vseg/nmwqptw8 during createAgent
To run a sweep, wrap train_model in a closure that binds dataset_file + node_name.
5. Writing NIFTI files#
The following script assembles all the predictions of one volume (inside a folder) into a nifti file. Since predict_one is used, all the individual segmentations are saved in the designated path. To do this, every segmentation is added to a list, then stacked as a 3D volume numpy array. The ants library saves it as a NIFTI file
[14]:
from liom_toolkit.segmentation.vseg.prediction import predict_one
from liom_toolkit.segmentation.vseg.model import VsegModel
device = "cuda" if torch.cuda.is_available() else "cpu"
model = VsegModel(pretrained=True, pretrained_artifact="liom-lab/model-registry/Vessel Segmentation:latest", device=device)
[15]:
import ants
import numpy as np
from tqdm.auto import tqdm
import os
import natsort
[16]:
# Prediction for every image in the directory
folder_dir = "data/vseg/S23_555nm_slices"
output_path = "data/vseg/LSFM_predictions"
os.makedirs(output_path, exist_ok=True)
normalization = True
image_list = os.listdir(folder_dir)
image_list = natsort.natsorted(image_list)
if image_list and image_list[-1] == '.ipynb_checkpoints':
image_list.remove('.ipynb_checkpoints')
image_3D = []
for images in tqdm(image_list):
image_path = f"{folder_dir}/{images}"
prediction = predict_one(model=model, img_path=image_path, save_path=output_path, norm=normalization, dev=device,
patching=False)
image_3D.append(prediction)
[17]:
volume = np.stack(image_3D)
volume = np.transpose(volume, (2, 1, 0))
[18]:
nifti = ants.from_numpy(volume)
nifti.to_file('data/vseg/brainslices.nii')
print("Saved data/vseg/brainslices.nii")
Saved data/vseg/brainslices.nii
6. Validation#
To test the model on new data.
The image paths are put into a list.
The masks named “XXXX_mask.png” should be in the same folder as each image.
The images should be in a folder with the volume name
The normalization is on and the patching is off
The following metrics are calculated:
Accuracy
Recall
Jaccard index
f1 score
In the save_path folder, the following files are saved:
The predictions -> “XXXX_segmented.png”
A comparison of the prediction and the mask -> “volume_XXXX_comparison.png”
In this image, the pixels are shown as follows:
TP: white
TN: black
FN: blue
FP: red
A CSV file with the metrics for each image and an average -> “validationmetrics.csv”
[19]:
from liom_toolkit.segmentation.vseg.validation import validate_model
images = ["data/vseg/s23/750.png", "data/vseg/s24/1200.png"]
save_path = "data/vseg/validation"
os.makedirs(save_path, exist_ok=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
validate_model(model=model, img_list=images, save_path=save_path, device=device)
7. Others#
Other code left over from the VSEG project
[20]:
from skimage.morphology import binary_erosion, disk
from skimage.io import imread, imsave
import numpy as np
from skimage.exposure import equalize_adapthist
from skimage.color import gray2rgb
[21]:
# Erodes a mask by one pixel
path = "data/vseg/s23/800_mask.png"
output_path = "data/vseg/eroded/800_mask.png"
os.makedirs(os.path.dirname(output_path), exist_ok=True)
mask = imread(path)
mask = (mask / mask.max()).astype(np.uint8)
erosion_disk = disk(1)
image = binary_erosion(mask, footprint=erosion_disk)
image = image.astype(np.uint8) * 255
imsave(output_path, image, cmap="gray")
/tmp/ipykernel_237921/1239335050.py:10: FutureWarning: `binary_erosion` is deprecated since version 0.26 and will be removed in version 0.28. Use `skimage.morphology.erosion` instead. Note the pixel shift by 1 for even-sized footprints (see docstring notes).
image = binary_erosion(mask, footprint=erosion_disk)
/tmp/ipykernel_237921/1239335050.py:12: FutureWarning: The plugin infrastructure in `skimage.io` is deprecated since version 0.25 and will be removed in 0.27 (or later). To avoid this warning, please do not pass additional keyword arguments for plugins (`**plugin_args`). Instead, use `imageio` or other I/O packages directly. See also `skimage.io.imsave`.
imsave(output_path, image, cmap="gray")
[22]:
# Converts the png values from 0-255 to 0-1
path = "data/vseg/s23/800_mask.png"
image = imread(path)
image = image / image.max()
image = image.astype(np.uint8)
imsave(path, image, check_contrast=False)
[23]:
# Saves the Clahe of the selected image
number = 500
volume = "S24"
image_path = f"data/vseg/{volume}_555nm_slices/{number}.png"
output_path = f"data/vseg/clahe/{volume}_{number}_eqhist.png"
os.makedirs(os.path.dirname(output_path), exist_ok=True)
image = imread(image_path)
image = (image / image.max() * 255).astype(np.uint8)
image_clahe = equalize_adapthist(image, kernel_size=10, clip_limit=0.05, nbins=128)
image_clahe = gray2rgb(image_clahe)
image_clahe = (image_clahe / image_clahe.max() * 255).astype(np.uint8)
imsave(output_path, image_clahe, cmap="gray")
/tmp/ipykernel_237921/2943563183.py:14: FutureWarning: The plugin infrastructure in `skimage.io` is deprecated since version 0.25 and will be removed in 0.27 (or later). To avoid this warning, please do not pass additional keyword arguments for plugins (`**plugin_args`). Instead, use `imageio` or other I/O packages directly. See also `skimage.io.imsave`.
imsave(output_path, image_clahe, cmap="gray")
[ ]: