Multiple channels#

Omnipose inherits the capability of Cellpose to segment based on multi-channel images. We will use this as an opportunity to show how we can run several models at once on the same image(s), in this case comparing Omnipose to Cellpose trained on the cyto2 dataset.

 1# First, import dependencies.
 2import numpy as np
 3import time, os, sys
 4from cellpose_omni import models, core, utils
 5
 6
 7# This checks to see if you have set up your GPU properly.
 8# CPU performance is a lot slower, but not a problem if 
 9# you are only processing a few images.
10use_GPU = core.use_gpu()
11print('>>> GPU activated? %d'%use_GPU)
12
13# for plotting 
14import matplotlib.pyplot as plt
15plt.style.use('dark_background')
16import matplotlib as mpl
17%matplotlib inline
18mpl.rcParams['figure.dpi'] = 300
19from omnipose.plot import imshow, colorize
Hide code cell output
2024-01-18 00:27:35,995 [INFO] ** TORCH GPU version installed and working. **
>>> GPU activated? 1

Load file#

This is one of the images from the cyto2 test dataset. Note that it is a good idea to always work with lists, even when the list of images is 1 long. It allows you to reuse your code easily when you do have a larger set of images to process.

 1from urllib.parse import urlparse
 2import skimage.io
 3
 4
 5urls = ['http://www.cellpose.org/static/images/img02.png']
 6files = []
 7for url in urls:
 8    parts = urlparse(url)
 9    filename = os.path.basename(parts.path)
10    if not os.path.exists(filename):
11        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, filename))
12        utils.download_url_to_file(url, filename)
13    files.append(filename)
14    
15imgs = [skimage.io.imread(f) for f in files]
16# print(imgs[0].shape)
17imgs = [np.stack((im[...,-1],im[...,1])) for im in imgs] # put cytosol in 1st channel, nucleus in 2nd
18nimg = len(imgs)
Hide code cell output
(349, 467, 3)

Read in the images from the file list. It's a good idea to display the images before proceeding.

 1from cellpose_omni import io, transforms 
 2
 3# print some infor about the images 
 4for i in imgs:
 5    print('img shape:',i.shape)
 6nimg = len(imgs)
 7# print(nimg)
 8
 9colors = np.array([[1,0,0],[0,1,0]])*1.0
10for k in range(len(imgs)):
11    imgs[k] = transforms.normalize99(imgs[k],omni=True)
12    rgb = colorize(imgs[k],colors=colors)    
13    imshow(rgb) 
img shape: (2, 349, 467)
../_images/47dcae9e3a9d4a5bdcbd0a9f3ebae287fc102bdd3d6a7551b68dd0605b4502c9.png

Initialize models#

We will compare two models here: cyto2 (Cellpose) and cyto2_omni. The latter was trained via the following command:

python -m cellpose --train --use_gpu --dir /home/kcutler/DataDrive/cyto2/train --mask_filter _masks --n_epochs 4000 --pretrained_model None  --learning_rate 0.1 --diameter 36 --save_every 50 --save_each --omni --verbose --chan 1 --chan2 2 --RAdam --batch_size 16 --img_filter _img
1model_name = ['cyto2','cyto2_omni']
2L = len(model_name)
3model = [models.CellposeModel(gpu=use_GPU, model_type=m) for m in model_name]
2024-01-18 00:31:08,507 [INFO] >>cyto2<< model set to be used
2024-01-18 00:31:08,540 [INFO] ** TORCH GPU version installed and working. **
2024-01-18 00:31:08,540 [INFO] >>>> using GPU
2024-01-18 00:31:08,651 [INFO] >>cyto2_omni<< model set to be used
2024-01-18 00:31:08,652 [INFO] ** TORCH GPU version installed and working. **
2024-01-18 00:31:08,652 [INFO] >>>> using GPU

Run segmentation#

The channels input can be very confusing. In the Cellpose documentation, it is stated that the list [chan,chan2] should represent the main channel to segment (chan) and the optional nuclear channel (chan2). But to train via CLI, chan is the "channel to segment" and chan2 is the nuclear channel, and the Cellpose team states the CLI command used to train their model used --chan 2 --chan2 1. Because 0 is grayscale and 1,2,3 are R,G,B (note the 1-based indexing here) this means that the given training command actually trains with G cytosol and R nuclei. This might imply that the cyto2_omni model actually is trained 'incorrectly', if I understood this right.

On top of this, the downloaded image has nuclei in channel 2 and cytosol in channel 1 (blue and green, respectively), whereas the cyto2 dataset shows cytosol as channel 0 and nuclei as channel 1 (this may be a result of using OpenCV, which usies BGR by default instead of RGB). So in fact, I should have trained the cyto2_omni model with --chan 1 --chan2 2. (Have not yet done this with most recent models...) Keep this in mind as you train your own models.

Note

You can train Omnipose on arbitrary numbers of channels using the --all_channels parameter. The reason for the cyto* models being trained with two channel was to focus on eukaryotes with cytosol and nuclei tags while also including many other images that might be single-channel. The chan, chan2 arguments and corresponding code in Cellpose is thus highly specific to this use case. Omnipose should probably always be trained with --all_channels on a dataset with the same number of channels across all images. After training, initialize the model with nchan=nchan and evaluate using chans=None.

For now, the following shows what channel arguments you need for the provided cyto2 models:

 1chans = [[2,1],[1,2]] # green cytoplasm [2] and red nucleus [1], see above 
 2n = range(nimg) 
 3
 4# define parameters
 5mask_threshold = [-1,-1,-1] #new model might need a bit lower 
 6verbose =0 # turn on if you want to see more output 
 7use_gpu = use_GPU #defined above
 8transparency = True # transparency in flow output
 9rescale= None # give this a number if you need to upscale or downscale your images
10flow_threshold = 0 # default is .4, but only needed if there are spurious masks to clean up; slows down output
11resample = False #whether or not to run dynamics on rescaled grid or original grid 
12
13N = L+1 # three options: pure cellpose, mixed, omnipose, new omnipose
14omni = [0,1,1]
15ind = [0,0,1]
16masks, flows, styles = [[]]*N, [[]]*N, [[]]*N
17
18diameter = 30
19for i in range(N):
20    masks[i], flows[i], styles[i] = model[ind[i]].eval([imgs[i] for i in n],channels=chans[ind[i]],
21                                                       diameter=diameter,
22                                                       mask_threshold=mask_threshold[i],
23                                                       transparency=transparency,
24                                                       flow_threshold=flow_threshold,
25                                                       omni=omni[i], #toggle omni
26                                                       resample=resample,verbose=verbose, 
27                                                       cluster=omni[i],
28                                                       interp=True, tile=False)

Plot the results#

 1from cellpose_omni import plot
 2import omnipose
 3
 4for idx,i in enumerate(n):
 5    
 6    for k,ki in enumerate(ind):
 7        
 8        print('model: {}, omni: {}'.format(model_name[ki],omni[ki]))
 9        maski = masks[k][idx] # get masks
10        flowi = flows[k][idx][0] # get RGB flows
11        imgi = omnipose.utils.normalize99(imgs[i])
12
13        # set up the output figure to better match the resolution of the images 
14        f = 10
15        szX = maski.shape[-1]/mpl.rcParams['figure.dpi']*f
16        szY = maski.shape[-2]/mpl.rcParams['figure.dpi']*f
17        fig = plt.figure(figsize=(szY,szX*4))
18        fig.patch.set_facecolor([0]*4)
19
20        plot.show_segmentation(fig, 
21                               imgi,
22                               maski, flowi, 
23                               channels=chans[i], 
24                               channel_axis=0,
25                               omni=True,
26                               img_colors=colors,
27                       interpolation=None)
28
29        plt.tight_layout()
30        plt.show()
model: cyto2, omni: 0
../_images/d3434576caa9e9c2ba9a5987dd890cc8c02de8d5213301707c0238ec79706faa.png
model: cyto2, omni: 0
../_images/2d9cebd9b68a9ccbcf315e3c31cd39d5cd31416239e89742b1181e6d49df4c4f.png
model: cyto2_omni, omni: 1
../_images/f4d12f22de138543bdee6c7e7afeac17dbcf8d22ee9b10c52bf604c1bfcf4fd9.png

Some comments on the above: Omnipose pre-processes the images slightly differently (see normalize99) and therefore the flow is a bit different even with the same model and input image compared to stock Cellpose. The cluster option helps a lot to get accurate masks with Omnipose in thin regions, but can result in under-segmentation between cells with poorly-defined flow fields. This can be a weakness of Omnipose relative to Cellpose, but as seen in the paper, Omnipose does slightly better than Cellpose on the cyto2 dataset on average. On roundish and low-accuracy datasets like cyto2, Omnipose simply does better in some areas and worse in others.