Fine-tune Cellpose on your own labels, with no GPU of your own

What you will do

The nuclei segmentation tutorial ends with a COCO export of your polygons. This page takes that export and fine-tunes Cellpose on it. You will finish with a model file and a number that says whether it does better than the stock model on your images.

You need a Google account, a COCO export from DataTorch, and the images you labeled. No GPU and no install on your own machine.

Fine-tuning means starting from Cellpose's built-in model and adjusting it with your labels. It is the route the Cellpose documentation recommends, and it needs far fewer images than training from nothing.

Step 1: put your data in Google Drive

In your Google Drive, make a folder named nuclei with this inside:

nuclei/
  export.json     your COCO export from DataTorch
  images/         the images you labeled, with the same file names

Step 2: open a notebook with a GPU

1. 1

   Create a notebook

   Go to colab.research.google.com and press New notebook.

2. 2

   Turn on the GPU

   Open Runtime, then Change runtime type, choose a GPU and press Save.

3. 3

   Connect your Drive and install Cellpose

   Run the cell below. Colab asks for permission to read your Drive.

from google.colab import drive
drive.mount("/content/drive")

!pip install cellpose tifffile pycocotools

Step 3: convert the export to Cellpose's format

Cellpose trains from pairs of files in one folder: an image, and a mask file with the same name plus _masks. In a mask, background is 0 and each nucleus has its own number.

This cell makes those pairs and holds back one image in five as a test set, which the model never trains on:

import random, shutil
from pathlib import Path

import numpy as np
import tifffile
from pycocotools.coco import COCO

root = Path("/content/drive/MyDrive/nuclei")
train_dir, test_dir = Path("/content/train"), Path("/content/test")
train_dir.mkdir(exist_ok=True)
test_dir.mkdir(exist_ok=True)

coco = COCO(str(root / "export.json"))
image_ids = coco.getImgIds()
random.seed(0)
random.shuffle(image_ids)
n_test = max(1, len(image_ids) // 5)

for n, image_id in enumerate(image_ids):
    info = coco.loadImgs(image_id)[0]
    name = Path(info["file_name"]).name
    out = test_dir if n < n_test else train_dir

    labels = np.zeros((info["height"], info["width"]), dtype=np.uint16)
    anns = coco.loadAnns(coco.getAnnIds(imgIds=image_id))
    for i, ann in enumerate(anns, start=1):
        labels[coco.annToMask(ann) > 0] = i

    shutil.copy(root / "images" / name, out / name)
    tifffile.imwrite(out / f"{Path(name).stem}_masks.tif", labels)

print(len(image_ids) - n_test, "training images,", n_test, "test images")

Step 4: fine-tune

from cellpose import io, models, train

io.logger_setup()
output = io.load_train_test_data(
    str(train_dir), str(test_dir), mask_filter="_masks", look_one_level_down=False
)
images, labels, image_names, test_images, test_labels, test_names = output

model = models.CellposeModel(gpu=True)
model_path, train_losses, test_losses = train.train_seg(
    model.net,
    train_data=images, train_labels=labels,
    test_data=test_images, test_labels=test_labels,
    weight_decay=0.1, learning_rate=1e-5, n_epochs=100,
    model_name="nuclei_v1",
)
shutil.copy(model_path, root / "nuclei_v1")
print(train_losses[-1], test_losses[test_losses > 0][-1])

These are the defaults in the Cellpose documentation. The first run downloads Cellpose's built-in model, about 1.2 GB. Expect about an hour for 40 images on the T4 that free Colab hands out, and watch the log: it prints the losses every ten epochs. Cellpose scores the test set only every five epochs, which is why the last line picks the last non-zero test loss. The shutil.copy line copies the model to your Drive, because Colab deletes its machine when the session ends.

Step 5: check that it helped

A lower loss does not tell you the model is better on your images. Score both models on the images you held back:

from cellpose import metrics

stock = models.CellposeModel(gpu=True)
tuned = models.CellposeModel(gpu=True, pretrained_model=str(model_path))

for name, m in [("stock", stock), ("fine-tuned", tuned)]:
    masks = m.eval(test_images)[0]
    ap = metrics.average_precision(test_labels, masks)[0]
    print(name, "AP at IoU 0.5:", ap[:, 0].mean())

Average precision at IoU 0.5 counts a nucleus as found when the model's outline overlaps the true one by at least half. Closer to 1 is better.

Our run. We did this with 50 images from the DSB2018 nuclei project, a mix of fluorescence, brightfield and stained tissue, split into 40 for training and 10 held back. Free Colab gave us a Tesla T4, and the 100 epochs of Step 4 took 59 minutes. The train loss ended at 0.147 and the test loss at 0.103. Step 5 printed:

stock AP at IoU 0.5: 0.837
fine-tuned AP at IoU 0.5: 0.858

A small gain from 40 images, and not the same on every image: the score rose on seven of the ten, was unchanged on one and fell a little on two. Treat it as one run, not a general result.

If the fine-tuned model does not beat the stock one, use the stock one. Your labels may be too few, or your images may already be easy for it.

If something goes wrong

• Colab gives no GPU. The free version does not always have one. Try again later. Training on a CPU is possible, but slow enough that it is rarely worth it.
• Out of memory. Crop or shrink very large images before Step 3.
• Images are skipped. Images with fewer than 5 masks are left out of training by default.
• A file name is not found in Step 3. The names in export.json must match the files in images/.

Next: put it behind a button

Your model file is ready to use on real work. Run it from a button in the annotator, correct what it gets wrong, and those corrections are your next training set.