AI Techniques

The chapter covers AI techniques for microscopic data at the example of E. coli imaging.

# Install necessary packages

# %pip install scikit-image
# %pip install opencv-python
# %pip install torch
# %pip install albumentations
# %pip install segmentation-models-pytorch
# %pip install pytorch-lightning
# %pip install cellpose

Challenging Images

We further increase the challenge to images with more bacteria.

For that purpose we will use the dataset by Scherr et al.. Please download it and move it into the main folder, so that you see the folder microbeSEG_dataset next to the notebook.

Let us view a (normalized) example image from the test dataset.

import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import skimage
from skimage.draw import disk
from scipy.ndimage import binary_erosion, binary_dilation
from scipy import ndimage as ndi

from skimage.segmentation import watershed
from skimage.feature import peak_local_max


def find_watershed(
        img: np.ndarray,
        min_distance: int=10,
        footprint: np.ndarray=np.ones((10,10))
    ) -> tuple[np.ndarray, int]:
    """Find watershed segmentation of an image.
    
    Arguments:
        img: input image (2D array)
        min_distance: minimum distance between local maxima
        footprint: search area for local maxima
    
    Returns:
        labels: segmentation labels of pixels
        num_features: number of detected features
    """
    
    # Compute distance to background (background has value 0)
    distance: np.ndarray = ndi.distance_transform_edt(img)
    assert isinstance(distance, np.ndarray), "ensure type"

    # Get local (within footprint) maxima of the distance field
    # filter via minimum distance between maxima
    coords = peak_local_max(distance, labels=img, min_distance=min_distance, footprint=footprint)

    # Create segmentation labels for pixels
    mask = np.zeros(distance.shape, dtype=bool)
    mask[tuple(coords.T)] = True
    markers, num_features = ndi.label(mask)
    return watershed(-distance, markers, mask=img), num_features


# Dataset path
dataset_path = Path("microbeSEG_dataset") / "30min-man_15min-pre"

# Load an example image
file_path = dataset_path / "train" / "img_000.tif"
img = skimage.io.imread(file_path)

# Normalize to range 0-255
img = img / np.max(img) * 255

plt.imshow(img, cmap=plt.cm.gray);

img_bin = np.zeros(img.shape).astype("int64")
mask = (img < 50)
img_bin[mask] = 100

img_bin = binary_erosion(img_bin, iterations=2)
img_bin = binary_dilation(img_bin, iterations=2)
plt.imshow(img_bin, cmap=plt.cm.gray);

On the right center of the image, it is already visible that some of the background has been detected as objects. You can play around with the preprocessing parameters and the threshold of 50, but no satisfying parameters can be found.

Indeed, subsequent segmentation with watershedding yields wrong results.

labels, num_features = find_watershed(img_bin, min_distance=10, footprint=np.ones([10,10]))
plt.imshow(labels, cmap=plt.cm.nipy_spectral)
print("Number of cells:", num_features)

Number of cells: 97

The number of cells has clearly been overestimated. Playing round with the parameters of find_watershed may slightly improve the result, but the ground truth is not (easily) obtained.

While including different pre-processing steps, such as normalizing and segmenting on local crops, may solve the problem of a global threshold and yield the correct result, finding the right pipeline can be challenging and time consuming. We explore a different approach in the following section.

Deep learning for segmentation

Neural networks (NNs) have shown impressive performance on multiple problems. We will use one to demonstrate the use of NNs for segmenting microscopic images. While several libraries for bacteria / cell segmentation exist, we will use a more general segmentation library in the following. The purpose of this is twofold:

• Currently available specialized libraries are limited in several aspects, making it difficult to adapt them to our needs. Typically one tries them all and picks the best for one’s need.
• They come as black boxes, and using them provides little insight into which algorithmic steps are used.
• Several of the libraries are outdated, not running on up-to-date Python libraries.

The classical approach to training NNs to solve segmentation is via supervised learning, training the NN with tuples of an input image img and a corresponding mask mask. This requires the availability of labeled masks. Luckily the dataset by Scherr et al. comes with such masks.

While coding your own Pytorch NNs is clearly possible, we skip this step and use the library segmentation_models which comes with a set of standard NN architectures for segmentation as well as the possibility to download pre-trained weights.

For example, a model with the Unet architecture is simply created via:

model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=OUT_CLASSES,            # model output channels (number of classes in your dataset)
)

import segmentation_models_pytorch as smp

OUT_CLASSES = 1

/Users/mfuegger/Github/Biodisco/computational_bioengineering_tuwien2026/venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Here, we will use a different architecture. We will train our model on the data within microbeSEG_dataset/30min-man_15min-pre/train/. The following code was adapted from an example segmentation of cars.

# Code adapted from https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb

import os
import cv2
import re

import torch
import numpy as np
import albumentations as A
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset
from torch.optim import lr_scheduler
import segmentation_models_pytorch as smp
import pytorch_lightning as pl

DATA_DIR = Path("microbeSEG_dataset") / "30min-man_15min-pre"
TRAIN_DIR = DATA_DIR / "train"
VAL_DIR = DATA_DIR / "val"
TEST_DIR = DATA_DIR / "test"

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[22], line 4
      1 # Code adapted from https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb
      2 
      3 import os
----> 4 import cv2
      5 import re
      6 
      7 import torch

File ~/Github/Biodisco/computational_bioengineering_tuwien2026/venv/lib/python3.12/site-packages/cv2/__init__.py:181
    176             if DEBUG: print("Extra Python code for", submodule, "is loaded")
    178     if DEBUG: print('OpenCV loader: DONE')
--> 181 bootstrap()

File ~/Github/Biodisco/computational_bioengineering_tuwien2026/venv/lib/python3.12/site-packages/cv2/__init__.py:153, in bootstrap()
    149 if DEBUG: print("Relink everything from native cv2 module to cv2 package")
    151 py_module = sys.modules.pop("cv2")
--> 153 native_module = importlib.import_module("cv2")
    155 sys.modules["cv2"] = py_module
    156 setattr(py_module, "_native", native_module)

File ~/.pyenv/versions/3.12.11/lib/python3.12/importlib/__init__.py:90, in import_module(name, package)
     88             break
     89         level += 1
---> 90 return _bootstrap._gcd_import(name[level:], package, level)

KeyboardInterrupt: 

The files have names composed of a prefix, an intermediate _, and a number with leading 0s. Let’s code a function that returns the images with a certain prefix. We need this to seperate img from mask files; and ensure they are ordered by increasing numbers.

def prefix_files(dir: str | Path, prefix: str) -> list[str]:
    """Get list of files in directory with given prefix, sorted by number suffix."""
    pattern = prefix + r'_(\d+)'

    # Extract filename and number
    img_files: list[tuple[str, int]] = []
    
    for filename in os.listdir(dir):
        match = re.match(pattern, filename)
        if match:
            img_files.append((filename, int(match.group(1))))
    
    # sort by number
    img_files.sort(key=lambda x: x[1])
    return [os.path.join(dir, f[0]) for f in img_files]

prefix_files(TRAIN_DIR, "img")

['microbeSEG_dataset/30min-man_15min-pre/train/img_000.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_001.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_004.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_005.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_008.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_009.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_011.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_012.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_014.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_015.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_017.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_020.tif',
 'microbeSEG_dataset/30min-man_15min-pre/train/img_021.tif']

We next code a class that allows us to retrieve data from an image directory. The class has some subtle features that will be import for training an NN-based segmenter:

• It extracts an image and returns a classification according to CLASSES per pixel. We can specify with a parameter classes which classes an existing mask file contains.
• It provides the possibility to apply an augmentation to the stored images and masks. Augmentations are image transformations. The augmentation is specified with augmentation.

Further, we code a simple image/mask visualization function.

# Code adapted from https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb

class Dataset(BaseDataset):
    """Read images, apply augmentation transformations.

    Args:
        images_dir (str): path to images folder
        class_values (list): values of classes to extract from segmentation mask
        augmentation (albumentations.Compose): data transformation pipeline
            (e.g. flip, scale, etc.)

    """

    CLASSES = [
        "cell",
        "background",
    ]

    def __init__(
        self,
        images_dir: Path | str,
        classes: tuple | list=(),
        augmentation=None,
        augmentation_factor: int=10,
    ):
        """Create the dataset."""
        self.images_fps = prefix_files(images_dir, "img") * augmentation_factor
        self.masks_fps = prefix_files(images_dir, "mask") * augmentation_factor

        # convert str names to class values on masks
        # e.g. ['cell'] -> [1]
        self.class_values = [self.CLASSES.index(c.lower()) for c in classes]

        self.augmentation = augmentation

    def __getitem__(self, i):
        """Read image and mask."""

        image = cv2.imread(self.images_fps[i])
        mask = cv2.imread(self.masks_fps[i], flags=2)

        assert image is not None, f"Image not found {self.images_fps[i]}"
        assert mask is not None, f"Mask not found {self.masks_fps[i]}"

        # Extract certain classes from mask (e.g. cells)
        masks = [(mask == v) for v in self.class_values]
        mask = np.stack(masks, axis=-1).astype("float32")
        
        # Invert mask: 1 = segment, 0 = background
        mask = 1 - mask

        # Apply augmentation to image and mask
        if self.augmentation is not None:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample["image"], sample["mask"]

        # Change data format from HWC (height-width-channel) to CHW (channel-height-width)
        # as required by PyTorch
        return image.transpose(2, 0, 1), mask.transpose(2, 0, 1)

    def __len__(self) -> int:
        """Number of images in the dataset, including augmentations."""
        return len(self.images_fps)
    

def visualize(**image_masks) -> None:
    """Plot images in one row."""
    n: int = len(image_masks)
    plt.figure(figsize=(7, 4))
    for i, (name, image) in enumerate(image_masks.items()):
        plt.subplot(1, n, i + 1)
        plt.title(" ".join(name.split("_")))
        
        if isinstance(image, torch.Tensor):
            image = image.cpu().numpy()

        if name == "image":
            plt.imshow(image.transpose(1, 2, 0))
        else:
            plt.imshow(image.squeeze())
        plt.axis("off")
    plt.show()

Let’s inspect one of the image/mask pairs and show it without any augmentation.

dataset = Dataset(TRAIN_DIR, classes=["cell"])

# Show one of them
image, mask = dataset[0]
visualize(image=image, cell_mask=mask)

Since there are relatively few images in the dataset, we will augment our dataset by applying transformations. This also ensures that a certain generalizability of the trained model to different environmental conditions.

Again, let’s inspect some the image/mask pairs.

# Code adapted from https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb

# training set images augmentation
def get_training_augmentation():
    train_transform = [
        A.HorizontalFlip(p=0.5),
        A.Affine(
            scale=(0.8, 1.2), rotate=(-45, 45), translate_percent=(-0.05, 0.05), p=1, border_mode=0
        ),
        A.PadIfNeeded(min_height=320, min_width=320, p=1),
        A.RandomCrop(height=320, width=320, p=1),
        A.GaussNoise(p=0.01),
        A.OneOf(
            [
                A.CLAHE(p=1),
                A.RandomBrightnessContrast(p=1),
                A.RandomGamma(p=1),
            ],
            p=0.9,
        ),
        A.OneOf(
            [
                A.Sharpen(p=1),
                A.Blur(blur_limit=3, p=1),
                A.MotionBlur(blur_limit=3, p=1),
            ],
            p=0.9,
        ),
        A.OneOf(
            [
                A.RandomBrightnessContrast(p=1),
                A.HueSaturationValue(p=1),
            ],
            p=0.9,
        ),
    ]
    return A.Compose(train_transform)


def get_validation_augmentation():
    """Add paddings to make image shape divisible by 32"""
    test_transform = [
        A.PadIfNeeded(min_height=384, min_width=480, p=1),
    ]
    return A.Compose(test_transform)

# Visualize resulted augmented images and masks
augmented_dataset = Dataset(
    TRAIN_DIR,
    augmentation=get_training_augmentation(),
    classes=["cell"],
)

# Check 3 augmented samples for the same image
for i in range(3):
    image, mask = augmented_dataset[0]
    visualize(image=image, mask=mask)

We are now in the position to code the data loaders for training, validation, and testing.

CLASSES = ["cell"]

train_dataset = Dataset(
    TRAIN_DIR,
    classes=CLASSES,
    augmentation=get_training_augmentation(),
)

valid_dataset = Dataset(
    VAL_DIR,
    classes=CLASSES,
    augmentation=get_validation_augmentation(),
)

test_dataset = Dataset(
    TEST_DIR,
    classes=CLASSES,
    augmentation=get_validation_augmentation(),
)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)
valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=0);

print(f"Training dataset has size: {len(train_dataset)}")

Training dataset has size: 130

# Code adapted from https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/cars%20segmentation%20(camvid).ipynb

# Some training parameters
EPOCHS = 60 # 60
T_MAX = EPOCHS * len(train_loader)

class BacteriaModel(pl.LightningModule):
    def __init__(self, arch: str, encoder_name: str, in_channels: int, out_classes: int, nn_model=None, **kwargs):
        super().__init__()
        # create nn_model if none is provided
        if nn_model is None:
            self.nn_model = smp.create_model(
                arch,
                encoder_name=encoder_name,
                in_channels=in_channels,
                classes=out_classes,
                **kwargs,
            )
        else:
            self.nn_model = nn_model

        # preprocessing parameters for image
        params = smp.encoders.get_preprocessing_params(encoder_name)
        self.register_buffer("std", torch.tensor(params["std"]).view(1, 3, 1, 1))
        self.register_buffer("mean", torch.tensor(params["mean"]).view(1, 3, 1, 1))

        # for image segmentation dice loss could be the best first choice
        self.loss_fn = smp.losses.DiceLoss(smp.losses.BINARY_MODE, from_logits=True)

        # initialize step metrics
        self.training_step_outputs = []
        self.validation_step_outputs = []
        self.test_step_outputs = []

    def forward(self, image):
        # normalize image
        image = (image - self.mean) / self.std
        mask = self.nn_model(image)
        return mask

    def shared_step(self, batch, stage: str):
        image, mask = batch

        # Required shape: (batch_size, num_channels, height, width)
        # For grayscale:  (batch_size, 1, height, width)
        assert image.ndim == 4

        # Check that image dimensions are divisible by 32,
        # encoder and decoder connected by `skip connections` and usually encoder have 5 stages of
        # downsampling by factor 2 (2 ^ 5 = 32)
        h, w = image.shape[2:]
        assert h % 32 == 0 and w % 32 == 0

        assert mask.ndim == 4

        # Check that mask values in between 0 and 1, NOT 0 and 255 for binary segmentation
        assert mask.max() == 1.0 and mask.min() == 0

        logits_mask = self.forward(image)

        # Predicted mask contains logits, and loss_fn param `from_logits` is set to True
        loss = self.loss_fn(logits_mask, mask)

        # convert mask values to probabilities, then
        prob_mask = logits_mask.sigmoid()
        # map to 0, 1
        pred_mask = (prob_mask > 0.5).float()

        # Return metrics
        tp, fp, fn, tn = smp.metrics.get_stats(
            pred_mask.long(), mask.long(), mode="binary"
        )
        return {
            "loss": loss,
            "tp": tp,
            "fp": fp,
            "fn": fn,
            "tn": tn,
        }

    def shared_epoch_end(self, outputs, stage):
        # aggregate metrics
        loss = [x["loss"].item() for x in outputs]
        tp = torch.cat([x["tp"] for x in outputs])
        fp = torch.cat([x["fp"] for x in outputs])
        fn = torch.cat([x["fn"] for x in outputs])
        tn = torch.cat([x["tn"] for x in outputs])

        # Per image IoU
        per_image_iou = smp.metrics.iou_score(
            tp, fp, fn, tn, reduction="micro-imagewise"
        )

        # Dataset IoU
        dataset_iou = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro")
        metrics = {
            f"{stage}_per_image_iou": per_image_iou,
            f"{stage}_dataset_iou": dataset_iou,
        }

        # Show the metrics in the progress bar
        self.log_dict(metrics, prog_bar=True, enable_graph=True)

    def on_train_epoch_end(self) -> None:
        self.shared_epoch_end(self.training_step_outputs, "train")
        self.training_step_outputs.clear()


    def training_step(self, batch, batch_idx):
        train_loss_info = self.shared_step(batch, "train")
        self.training_step_outputs.append(train_loss_info)
        return train_loss_info

    def validation_step(self, batch, batch_idx):
        valid_loss_info = self.shared_step(batch, "valid")
        self.validation_step_outputs.append(valid_loss_info)
        return valid_loss_info

    def on_validation_epoch_end(self) -> None:
        self.shared_epoch_end(self.validation_step_outputs, "valid")
        self.validation_step_outputs.clear()

    def test_step(self, batch, batch_idx):
        test_loss_info = self.shared_step(batch, "test")
        self.test_step_outputs.append(test_loss_info)
        return test_loss_info

    def on_test_epoch_end(self) -> None:
        self.shared_epoch_end(self.test_step_outputs, "test")
        # empty set output list
        self.test_step_outputs.clear()

    def configure_optimizers(self):
        # optimizer = torch.optim.Adam(self.parameters(), lr=2e-4)
        optimizer = torch.optim.Adam(self.parameters(), lr=2e-4)
        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=T_MAX, eta_min=1e-5)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            },
        }

We will use an architecture that is called Feature Pyramid Network (FPN).

# check if a saved model exists
model_path = Path("bacterial_model.pth")
if model_path.exists():
    print("Loading saved model from disk.")
    nn_model = torch.load(model_path, weights_only=False)
    bacterial_model = BacteriaModel(
        arch="FPN",
        encoder_name="resnext50_32x4d",
        in_channels=3,
        out_classes=OUT_CLASSES,
        nn_model=nn_model
    )
    trainer = pl.Trainer(max_epochs=0, log_every_n_steps=1)
    trainer.fit(
        bacterial_model,
        train_dataloaders=train_loader,
        val_dataloaders=valid_loader,
    )
else:
    print("Starting with a fresh model.")
    bacterial_model = BacteriaModel(
        arch="FPN",
        encoder_name="resnext50_32x4d",
        in_channels=3,
        out_classes=OUT_CLASSES,
        nn_model=None
    )

    # train it
    trainer = pl.Trainer(max_epochs=EPOCHS, log_every_n_steps=1)
    trainer.fit(
        bacterial_model,
        train_dataloaders=train_loader,
        val_dataloaders=valid_loader,
    )
    # save trained model to disk
    torch.save(bacterial_model.nn_model, model_path)

💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py:76: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default
You are using a CUDA device ('NVIDIA GeForce RTX 3090') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision

Loading saved model from disk.

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

┏━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name     ┃ Type     ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ nn_model │ FPN      │ 25.6 M │ train │     0 │
│ 1 │ loss_fn  │ DiceLoss │      0 │ train │     0 │
└───┴──────────┴──────────┴────────┴───────┴───────┘

Trainable params: 25.6 M                                                                                           
Non-trainable params: 0                                                                                            
Total params: 25.6 M                                                                                               
Total estimated model params size (MB): 102                                                                        
Modules in train mode: 210                                                                                         
Modules in eval mode: 0                                                                                            
Total FLOPs: 0                                                                                                     

/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/rich/live.py:256: UserWarning:
install "ipywidgets" for Jupyter support
  warnings.warn('install "ipywidgets" for Jupyter support')

/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/pytorch_lightning/trainer/conn
ectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider 
increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.

/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/pytorch_lightning/trainer/conn
ectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider
increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.

`Trainer.fit` stopped: `max_epochs=0` reached.

And finally run validation and test.

# run validation dataset
valid_metrics = trainer.validate(bacterial_model, dataloaders=valid_loader, verbose=False)
print(valid_metrics)

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

[{'valid_per_image_iou': 0.6897271871566772, 'valid_dataset_iou': 0.7325970530509949}]

# run test dataset
test_metrics = trainer.test(bacterial_model, dataloaders=test_loader, verbose=False)
print(test_metrics)

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.

[{'test_per_image_iou': 0.6965533494949341, 'test_dataset_iou': 0.6922621130943298}]

Let’s check the results on some image/mask/prediction tuples.

# Predict on a batch of test images
def predict(images: torch.Tensor, model: BacteriaModel, min_volume: int = 0) -> torch.Tensor:
    with torch.no_grad():
        bacterial_model.eval()
        logits = bacterial_model(images)
    pr_masks = logits.sigmoid()
    pr_masks = (pr_masks > 0.5).float()

    if min_volume > 0:
        # Remove small connected components
        # Find connected components
        labeled_array, num_components = ndi.label(pr_masks.numpy())
        
        # Calculate size of each component
        component_sizes = np.bincount(labeled_array.ravel())
        
        # Create mask for components larger than min_volume
        small_components = component_sizes < min_volume
        small_components[0] = False  # Keep background (label 0)
        
        # Remove small components
        pr_masks[small_components[labeled_array]] = 0
        
        # Relabel remaining components as 1
        pr_masks = (pr_masks > 0).float()
    
    return pr_masks



images, masks = list(next(iter(test_loader)))[:2]
pr_masks = predict(images, bacterial_model, min_volume=50)

for i in range(len(pr_masks)):
    if i > 3:
        break
    visualize(image=images[i], ground_truth=masks[i], predicted=pr_masks[i])

Limitations of generalizability

Equipped with a trained model, let’s try to segment an image from a different setup and dataset: a full image that was taken in our lab from E. coli.

We code some normalization that we will be using, first.

def gray_norm(img: np.ndarray, max_value: float = 99) -> np.ndarray:
    """Convert to grayscale and normalize to 0-max_value."""
    if img.ndim == 3:
        img = rgb2gray(img).astype(np.float32)
    img_normalized = (img * max_value).astype(np.float32)
    return img_normalized

And are ready to run the model on it.

assert bacterial_model is not None, "train model first"

image = cv2.imread("data/ecoli_full.png")
transformation = A.CenterCrop(height=32*10, width=32*10, p=1.0)
image = transformation(image=image)["image"].transpose(2, 0, 1)

image_tensor = torch.tensor(image)
pr_mask = predict(image_tensor, bacterial_model, min_volume=0)

visualize(image=image, predicted=pr_mask)

The segmentation is clearly not good. The image was too different from the ones it was trained on.

Dedicated Cell Segmenters

In the previous section we segmented microscopy with a general NN-based architecture. The framework is not specific to cells and may as well be run on dogs or cars.

We will look into cellpose as an example for a dedicated cell segmenter. Cellpose is a specialized tool for segmenting bacteria and other cells.

# %pip install cellpose
# %pip install scipy
# %pip install napari

Let’s run it on the first of the E. coli images.

import cv2
import albumentations as A
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import scipy

# HACK START
# Monkey-patch missing attributes
import numpy.fft as _npfft
scipy.fft = _npfft.fft
scipy.ifft = _npfft.ifft
# HACK END

from cellpose import models
from skimage import io
# from skimage.color import rgb2gray

# Load the E. coli image
img_path = Path("data") / "ecoli" / "img_0000.tif"
image = cv2.imread(img_path)
transformation = A.CenterCrop(height=32*10, width=32*10, p=1.0)
image = transformation(image=image)["image"].astype(np.float32)
img_normalized = gray_norm(image, max_value=99)

# Segment
model = models.CellposeModel(gpu=True, pretrained_model='bact_phase_omni')
masks, flows, styles = model.eval(
    img_normalized,
)

def cell_segmenter_visualize(
        img: np.ndarray,
        mask: np.ndarray,
    ) -> None:
    """Visualize cell segmentation results."""
    print(f"Detected {mask.max()} cells")

    plt.figure(figsize=(8, 4))

    plt.subplot(1, 2, 1)
    plt.imshow(img, cmap='gray')
    plt.title("image")
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.imshow(mask, cmap='nipy_spectral')
    plt.title(f"mask")
    plt.axis('off')

    plt.tight_layout()
    plt.show()

# Visualize results
cell_segmenter_visualize(img_normalized, masks)

pretrained model /home/mfuegger/.cellpose/models/cpsam not found, using default model

Detected 59 cells

The result is great and disappointing at the same time. It found all the round circles, that are likely from the lense. There is no way it could have known that we are nto interested in these, but the oval bacterial cells.

from matplotlib.pyplot import imsave
import napari
from skimage.io import imsave


# Load the E. coli image
ecoli_path = Path("data") / "ecoli"

def annotate_ecoli_img(nr: int) -> None:
    """Annotate file img_{nr}.tif using Napari viewer."""

    # load image
    image = cv2.imread(ecoli_path / f"img_{nr:04d}.tif").astype(np.float32)
    img_normalized = gray_norm(image, max_value=99)

    # load mask layer
    mask_path = ecoli_path / f"mask_{nr:04d}.tif"
    if mask_path.exists():
        print("Loading existing mask.")
        labels_layer_data = io.imread(mask_path).astype(np.uint16)
    else:
        print("Please annotate from empty mask.")
        labels_layer_data = np.zeros(img_normalized.shape[:2], dtype=np.uint16)

    # start viewer
    viewer = napari.Viewer()
    viewer.add_image(img_normalized, name="Microscopy")
    viewer.add_labels(labels_layer_data, name='Ecoli');

    # Show the viewer and block execution until it is closed
    viewer.show(block=True)

    # Get the mask data from the viewer
    labels = viewer.layers['Ecoli'].data

    # save mask
    imsave(mask_path, labels.astype('uint16'))

    # Close the viewer
    viewer.close()

    # Visualize results
    cell_segmenter_visualize(img_normalized, labels)

Let’s use this code to view the already annotated image nr 0. When you are good with the shown labels, close it and it will save the label to the mask file nr 0.

# Uncomment to show and annotate the first image
# annotate_ecoli_img(0)

Let’s do the same with image nr 10.

# annotate_ecoli_img(10)

# Train on E. coli dataset

from cellpose import models
import numpy as np
# from scipy import ndimage as ndi
import cv2

# Prepare training data
train_images = []
train_masks = []

# Load training data from the dataset
train_dir = ecoli_path

# Get image and mask files
train_nrs = [0, 10]

img_files = [train_dir / f"img_{i:04d}.tif" for i in train_nrs]
mask_files = [train_dir / f"mask_{i:04d}.tif" for i in train_nrs]

# Setting a seed for reproducibility
seed = 420  # 420
transformation = A.Compose([
    # A.RandomRotate90(p=0.5),
    A.Rotate(limit=45, p=0.5),
    A.RandomCrop(height=32*5, width=32*5, p=1.0),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.01),
], seed=seed)


print(f"Loading {len(img_files)} training images...")
for img_path, mask_path in zip(img_files, mask_files):
    # Read image
    image = cv2.imread(img_path).astype(np.float32)
    img_normalized = gray_norm(image, max_value=99)

    # Read mask
    mask = cv2.imread(mask_path, flags=cv2.IMREAD_ANYDEPTH).astype(np.uint16)

    # Apply a random crop
    for _ in range(20):
        transformed = transformation(image=img_normalized, mask=mask)
        img_transformed = transformed["image"]
        mask_transformed = transformed["mask"]
        train_images.append(img_transformed)
        train_masks.append(mask_transformed)

print(f"With transformations: {len(train_images)} image-mask pairs.")

Loading 2 training images...
With transformations: 40 image-mask pairs.

Whenever we generate image training data, it is a good idea to take a look at some of this data. It is very easy to get a bitsize, normalization, or order of the channels wrong.

We also set a seed in the transformation, so that the training data is reproducible.

Let’s look at the 1st training pair.

cell_segmenter_visualize( train_images[0], train_masks[0] )

Detected 11 cells

All looks great, ready to setup the training!

# Initialize model
model = models.CellposeModel(gpu=True, model_type='bact_phase_omni')

model_type argument is not used in v4.0.1+. Ignoring this argument...

Let’s take a quick look at the architecture of the underlying NN. It is a Transformer, a widely-used architecture.

model.net

Transformer(
  (encoder): ImageEncoderViT(
    (patch_embed): PatchEmbed(
      (proj): Conv2d(3, 1024, kernel_size=(8, 8), stride=(8, 8))
    )
    (blocks): ModuleList(
      (0-23): 24 x Block(
        (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
        (attn): Attention(
          (qkv): Linear(in_features=1024, out_features=3072, bias=True)
          (proj): Linear(in_features=1024, out_features=1024, bias=True)
        )
        (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
        (mlp): MLPBlock(
          (lin1): Linear(in_features=1024, out_features=4096, bias=True)
          (lin2): Linear(in_features=4096, out_features=1024, bias=True)
          (act): GELU(approximate='none')
        )
      )
    )
    (neck): Sequential(
      (0): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): LayerNorm2d()
      (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (3): LayerNorm2d()
    )
  )
  (out): Conv2d(256, 192, kernel_size=(1, 1), stride=(1, 1))
)

We ready to start the training.

There is one remark to do here, though. We trained this about 500 epochs, which we had to stop on a Mac Pro after 1h and moved to an NVIDIA GPU machine from our group. This then took about 10min Alternatively, you can download a pre-trained model from:

https://huggingface.co/mfuegger/ecoli_model

You can do this by uncommenting the following cell.

# %pip install huggingface_hub

import os
from huggingface_hub import hf_hub_download
import shutil
from pathlib import Path

def download_hf_model(repo_id: str, filename: str, model_dir: str | Path = "models") -> None:
    """Download a model from Hugging Face Hub and save it locally."""
    model_dir = Path(model_dir)

    # mind: often symlinks as cached_path returned
    cached_path = Path(hf_hub_download(repo_id=repo_id, filename=filename))
    real_path = cached_path.resolve(strict=True)
    print("Downloaded into HF cache:", real_path)
    
    # Path relative to current working directory (this script)
    local_dir = Path().resolve() / model_dir
    local_dir.mkdir(parents=True, exist_ok=True)
    local_path = local_dir / filename
    print("Saving to:", local_path)
    shutil.copy2(real_path, local_path)


# -----------------------------------------------------
# Uncomment if you want to download a pretrained model
# -----------------------------------------------------

# download_hf_model(repo_id="mfuegger/ecoli_model", filename="ecoli_model.pt")

Before we start the training, let’s check the parameters we chose, here. For example, we set min_train_masks = 0 since we want to train even if 0 objects are present. Since E. coli are sparse on the images, we would drop many of them with the default minimum of 5.

# Train the model
# Attention: epochs are set to 0 for demo purposes, set to 100 for real training

from cellpose import models, train

model_path, train_losses, test_losses = train.train_seg(
    net = model.net,
    train_data = train_images,
    train_labels = train_masks,
    test_data = None,
    test_labels = None,
    learning_rate = 1e-5,
    n_epochs = 10,  # 1 is for demo, 100 for real training
    batch_size = 8,  # you can decrease to 4 if running out of memory in demo
    weight_decay = 0.0001,
    normalize = True,
    min_train_masks = 0,
    model_name = "ecoli_model.pt",  # within models/ folder
    save_each = True,  # save each as seperate file
    save_every = 5,  # epochs
)

100%|██████████| 40/40 [00:00<00:00, 184.50it/s]
  0%|          | 0/40 [00:00<?, ?it/s]/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/numpy/_core/fromnumeric.py:3860: RuntimeWarning: Mean of empty slice.
  return _methods._mean(a, axis=axis, dtype=dtype,
/home/mfuegger/Github/computational_bioengineering/venv/lib/python3.12/site-packages/numpy/_core/_methods.py:145: RuntimeWarning: invalid value encountered in scalar divide
  ret = ret.dtype.type(ret / rcount)
100%|██████████| 40/40 [00:00<00:00, 8462.66it/s]

Let us check its performance on an image it was trained on. While this is clearly not a test for it generalizing-capabilities, this is good practice to check if the training took into account the training data.

We will be plotting the predicted result and below the expected result.

# Load model
pretrained_model = Path("models") / "ecoli_model.pt"
model = models.CellposeModel(gpu=True, pretrained_model=str(pretrained_model))

# Evaluate on a training image
labels, flows, styles = model.eval(
    train_images[0],
)

print("Cellpose segmentation results:")
cell_segmenter_visualize(train_images[0], labels)

Cellpose segmentation results:
Detected 4 cells

print("Training data:")
cell_segmenter_visualize(train_images[0], train_masks[0])

Training data:
Detected 5 cells

The result is good, but still contains some initially detected cells for the dots on the image. This shows that we should train for much longer.

Finally, we are going to test the network’s generalizing capabilities by running it on a previously unseen image.

nr = 25  # image number to test
image = cv2.imread(ecoli_path / f"img_{nr:04d}.tif").astype(np.float32)

x_min = 300
y_min = 250
H = 32*5
W = 32*5

# Draw rectangle to show crop area
image_with_rect = image.copy()
cv2.rectangle(
    image_with_rect,
    (x_min, y_min),
    (x_min + 32*5, y_min + 32*5),
    (255, 0, 0),
    2,
)
plt.imshow(image_with_rect.astype(np.uint8))
plt.title("Crop area")

# Crop the image
transformation = A.Crop(y_min=y_min, y_max=y_min + 32*5, x_min=x_min, x_max=x_min + 32*5)
image_transformed = transformation(image=image)["image"].astype(np.float32)

# Segment the cropped image
img_normalized = gray_norm(image_transformed, max_value=99)
labels, flows, styles = model.eval(
    img_normalized,
)
cell_segmenter_visualize(img_normalized, labels)

Detected 2 cells

This looks great! Of course, a large part of the real work begins now:

• Checking the performance on some other pictures.
• Validating the performance, e.g., with IoU and cell counts systematically over a test set.
• Changing learning meta-parameters, to improve results.
• Final checking on a validation set.

---

License: © 2025 Matthias Függer and Thomas Nowak. Licensed under CC BY-NC-SA 4.0.