Cellpose trained to learn scale#
This example shows how to train cellpose to find objects at different scales. We try to use the same cellpose model to detect several spheres of varying sizes in the same image.
When applying a cellpose model to an image, it is common to use the ‘diameter’ parameter to rescale image to detect structure at a different scale than the network was trained for. However this may not work if single images have objects at different scales within the image.
If the training data has a mixture of small and large objects, cellpose can learn these. This can be shown with a simple simulation.
import raster_geometry as rg
import numpy as np
from tnia.simulation.phantoms import add_small_to_large_2d
import matplotlib.pyplot as plt
from tnia.plotting.plt_helper import imshow_multi2d, imshow2d
import math
Make a super simple fake dataset#
Here we add a few different objects of different sizes.
To create a training set we a keep a copy of the single spheres (which are put in 224 by 224 arrays by raster geometry… perfect for training patches)
width, height = 624, 224
image = np.zeros([height, width], dtype=np.float32)
truth = np.zeros([height, width], dtype=np.float32)
rs = [3, 5, 15, 30, 60, 70]
x_ = 44
# for train will be a list of the 224 by 224 circle images that will be used for training
for_train = []
i = 0
for r in rs:
x, y = x_, 112
size = [math.ceil(r*2), math.ceil(r*2)]
size = [224, 224]
temp=rg.circle(size, r)
for_train.append(temp)
add_small_to_large_2d(image, temp, x, y, mode='replace_non_zero')
add_small_to_large_2d(truth, i*temp, x, y, mode='replace_non_zero')
x_ = x_ + 50+2*r
i += 1
fig = imshow2d(image, width = 8, height = 3.5)
# set title
fig.suptitle('Image')
fig = imshow2d(truth, width=8, height=3.5)
# set title
stop = fig.suptitle('Truth')
Load Cyto2#
We also print the mean diameter of cyto2 (if log is on this will print by default too)
from cellpose import models, io
model_cyto2 = models.CellposeModel(gpu=True, model_type="cyto2")
print('cyto2 diameter mean', model_cyto2.diam_labels)
2024-10-31 19:57:01,797 [INFO] >> cyto2 << model set to be used
2024-10-31 19:57:01,799 [INFO] ** TORCH CUDA version installed and working. **
2024-10-31 19:57:01,800 [INFO] >>>> using GPU
2024-10-31 19:57:01,941 [INFO] >>>> loading model C:\Users\bnort\.cellpose\models\cyto2torch_0
2024-10-31 19:57:02,119 [INFO] >>>> model diam_mean = 30.000 (ROIs rescaled to this size during training)
cyto2 diameter mean 30.0
Process with Cyto2 with different diameters#
We try a few different diameters to see if we can find all the objects. Unfortunately there is no diameter value that works for everything.
It looks like d=6 finds the small objects, d=140 finds the big one, and default finds only a medium one.
flow_threshold = 0
cellprob_threshold = -0.8
labels_d6 = model_cyto2.eval(image, channels=[0, 0],diameter=6,flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)[0]
labels_default = model_cyto2.eval(image, channels=[0, 0], flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)[0]
labels_none = model_cyto2.eval(image, channels=[0, 0], diameter = None, flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)[0]
labels_d70 = model_cyto2.eval(image, channels=[0, 0],diameter=70,flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)[0]
labels_d140 = model_cyto2.eval(image, channels=[0, 0],diameter=140,flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)[0]
fig = imshow_multi2d([image, labels_d6, labels_default, labels_none, labels_d70, labels_d140], ['truth', 'cyto2 d=6', 'cyto2 default', 'cyto2 None', 'cyto2 d=70','cyto2 d=140'], 6 ,1, width = 12, height = 22)
Repeat the objects to make a fake training set#
Just repeat a few times to make a trivial training set. No need to even add blur or noise, we just want to see if a custom model can learn different scales.
X = []
Y = []
for train in for_train:
X = X+[train.copy().astype('float32')[..., np.newaxis] for i in range(10)]
Y = Y+[train.copy().astype('uint16') for i in range(10)]
print(train.min(), train.max())
len(X)
False True
False True
False True
False True
False True
False True
60
Make a custom model#
import os
model_path = r'./'
model_name = 'custom'
logger = io.logger_setup()
model_custom = models.CellposeModel(gpu=True, model_type=None)#, pretrained_model=os.path.join(model_path,'models',model_name))
creating new log file
2024-10-31 19:45:43,152 [INFO] WRITING LOG OUTPUT TO C:\Users\bnort\.cellpose\run.log
2024-10-31 19:45:43,153 [INFO]
cellpose version: 3.0.9
platform: win32
python version: 3.10.14
torch version: 2.2.2+cu118
2024-10-31 19:45:43,156 [INFO] ** TORCH CUDA version installed and working. **
2024-10-31 19:45:43,157 [INFO] >>>> using GPU
2024-10-31 19:45:43,274 [INFO] >>>> no model weights loaded
Train our custom model#
Note our fake training set was trivial, a set of images each containing single circle ranging from large circles to small.
However I found that in order to learn this very simple data, we need to explicitly set ‘rescale’ to False when training.
Otherwise the images are rescaled during the training process, and we don’t actually learn the different scales.
from cellpose import train
test=train.train_seg(model_custom.net, X, Y,
channels=[0,0],
save_path=model_path,
n_epochs=180,
min_train_masks=1,
normalize = False,
rescale = False,
model_name=model_name)
2024-10-31 19:45:51,212 [INFO] computing flows for labels
100%|██████████| 60/60 [00:03<00:00, 15.67it/s]
2024-10-31 19:45:55,063 [INFO] >>> computing diameters
100%|██████████| 60/60 [00:00<00:00, 3177.30it/s]
2024-10-31 19:45:55,088 [INFO] >>> using channels [0, 0]
2024-10-31 19:45:55,102 [INFO] >>> n_epochs=180, n_train=60, n_test=None
2024-10-31 19:45:55,103 [INFO] >>> AdamW, learning_rate=0.00500, weight_decay=0.00001
c:\Users\bnort\miniconda3\envs\pytorch_and_SAM3\lib\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
2024-10-31 19:45:56,518 [INFO] >>> saving model to models\custom
2024-10-31 19:45:58,628 [INFO] 0, train_loss=1.3979, test_loss=0.0000, LR=0.0000, time 2.11s
2024-10-31 19:46:02,912 [INFO] 5, train_loss=0.6603, test_loss=0.0000, LR=0.0028, time 6.40s
2024-10-31 19:46:07,598 [INFO] 10, train_loss=0.2332, test_loss=0.0000, LR=0.0050, time 11.08s
2024-10-31 19:46:16,200 [INFO] 20, train_loss=0.0675, test_loss=0.0000, LR=0.0050, time 19.68s
2024-10-31 19:46:24,772 [INFO] 30, train_loss=0.0446, test_loss=0.0000, LR=0.0050, time 28.26s
2024-10-31 19:46:33,404 [INFO] 40, train_loss=0.0205, test_loss=0.0000, LR=0.0050, time 36.89s
2024-10-31 19:46:42,490 [INFO] 50, train_loss=0.0246, test_loss=0.0000, LR=0.0050, time 45.97s
2024-10-31 19:46:51,256 [INFO] 60, train_loss=0.0340, test_loss=0.0000, LR=0.0050, time 54.74s
2024-10-31 19:47:00,675 [INFO] 70, train_loss=0.0157, test_loss=0.0000, LR=0.0050, time 64.16s
2024-10-31 19:47:09,308 [INFO] 80, train_loss=0.0114, test_loss=0.0000, LR=0.0050, time 72.79s
2024-10-31 19:47:18,246 [INFO] 90, train_loss=0.0096, test_loss=0.0000, LR=0.0050, time 81.73s
2024-10-31 19:47:27,333 [INFO] 100, train_loss=0.0107, test_loss=0.0000, LR=0.0050, time 90.82s
2024-10-31 19:47:37,041 [INFO] 110, train_loss=0.0097, test_loss=0.0000, LR=0.0050, time 100.52s
2024-10-31 19:47:46,038 [INFO] 120, train_loss=0.0059, test_loss=0.0000, LR=0.0050, time 109.52s
2024-10-31 19:47:55,011 [INFO] 130, train_loss=0.0088, test_loss=0.0000, LR=0.0025, time 118.49s
2024-10-31 19:48:04,208 [INFO] 140, train_loss=0.0054, test_loss=0.0000, LR=0.0006, time 127.69s
2024-10-31 19:48:12,957 [INFO] 150, train_loss=0.0034, test_loss=0.0000, LR=0.0002, time 136.44s
2024-10-31 19:48:21,825 [INFO] 160, train_loss=0.0037, test_loss=0.0000, LR=0.0000, time 145.31s
2024-10-31 19:48:30,924 [INFO] 170, train_loss=0.0034, test_loss=0.0000, LR=0.0000, time 154.41s
labels_custom, c, d = model_custom.eval(image, channels=[0, 0], normalize = False)
truth.sum(), labels_custom.sum()
(132112.0, 49928)
Show all the cyto labels and label custom#
Note: cyto can only find objects of different sizes by tuning the diameter param. Custom can find very different object sizes by default.
import random
from tnia.plotting.plt_helper import random_label_cmap
cmaps = [random_label_cmap()]*7
fig = imshow_multi2d([truth, labels_d6, labels_default, labels_none, labels_d70, labels_d140, labels_custom],
['truth', 'cyto2 d=6', 'cyto2 default', 'cyto2 None', 'cyto2 70', 'cyto2 d=140', 'custom'],7 ,1, width = 12, height = 26, colormaps=cmaps)
plt.figure(figsize=(10,10))
plt.imshow(c[0])
plt.title('flows')
Text(0.5, 1.0, 'flows')
print(model_custom.diam_labels)
print(model_cyto2.diam_labels)
30.0
30.0
import napari
viewer = napari.Viewer()
viewer.add_image(c[0])