1# -*- coding: utf-8 -*-
2"""
3Created on Mon May 15 16:08:02 2023
4
5@author: Bart Steemans. Govers Lab.
6"""
7
8from cellpose_omni import models, core, io
9from cellpose_omni import models
10from cellpose_omni.models import MODEL_NAMES
11from numbers import Integral
12
13
14class Omnipose:
15
16 def __init__(self, imgs, paths):
17
18 self.imgs = imgs
19 self.files = paths
20 self.use_GPU = core.use_gpu()
21 self.masks = None
22 self.flows = None
23 self.selected_imgs = []
24 self.selected_files = []
25 self.model = None
26 print(f">>> GPU activated? {self.use_GPU}")
27
28 def load_models(self, model_name):
29 """
30 Load the desired omnipose model.
31
32 """
33 if self.model is None:
34 self.model = models.CellposeModel(gpu=self.use_GPU, model_type=model_name)
35
36 def process(self, n=(0,), mask_thresh=1, minsize=300):
37 """
38 Segment the raw images and get their masks.
39
40 """
41 chans = [0, 0] # th is means segment based on first channel, no second channel
42 mask_threshold = mask_thresh # negative = mask becomes larger, positive = mask becomes thinner
43 verbose = 0 # turn on if you want to see more output
44 use_gpu = self.use_GPU # defined above
45 transparency = True # transparency in flow output
46 rescale = (
47 None # give this a number if you need to upscale or downscale your images
48 )
49 omni = True # we can turn off Omnipose mask reconstruction, not advised
50 flow_threshold = 0.0 # default is .4, but only needed if there are spurious masks to clean up; slows down output
51 resample = (
52 True # whether or not to run dynamics on rescaled grid or original grid
53 )
54 batch_size = 1
55 niter = None
56 # diamter: too small -> oversegmentation, too large -> undersegmentation
57 # if diameter = None, the diam will be estimated (in images from test-set )
58 self.masks = None
59 self.flows = None
60 self.selected_imgs = []
61 self.selected_files = []
62 if len(self.imgs) != len(self.files):
63 raise ValueError("Image and filename counts do not match")
64 indices = list(range(len(self.imgs))) if n is None else list(n)
65 if not indices:
66 raise ValueError("At least one image must be selected for segmentation")
67 if len(set(indices)) != len(indices) or any(
68 not isinstance(i, Integral) or isinstance(i, bool)
69 or i < 0 or i >= len(self.imgs) for i in indices
70 ):
71 raise ValueError("Segmentation indices must be unique valid image positions")
72 self.selected_imgs = [self.imgs[i] for i in indices]
73 self.selected_files = [self.files[i] for i in indices]
74 self.masks, self.flows, styles = self.model.eval(
75 self.selected_imgs,
76 niter=niter,
77 channels=chans,
78 rescale=rescale,
79 mask_threshold=mask_threshold,
80 transparency=transparency,
81 flow_threshold=flow_threshold,
82 omni=omni,
83 resample=resample,
84 cluster=True,
85 verbose=verbose,
86 diameter=None,
87 min_size=minsize,
88 batch_size=batch_size,
89 )
90
91 def save_masks(self):
92 """
93 Save the segmentation masks.
94
95 """
96 if self.masks is None or self.flows is None:
97 raise RuntimeError("Run process() successfully before saving masks")
98 if len(self.masks) != len(self.selected_imgs) or len(self.flows) != len(self.selected_imgs):
99 raise ValueError("Segmentation results do not match selected images")
100 io.save_masks(
101 self.selected_imgs,
102 self.masks,
103 self.flows,
104 self.selected_files,
105 png=False,
106 tif=True, # whether to use PNG or TIF format
107 suffix="", # suffix to add to files if needed
108 save_flows=False, # saves both RGB depiction as *_flows.png and the raw components as *_dP.tif
109 save_outlines=False, # save outline images
110 dir_above=0, # save output in the image directory or in the directory above (at the level of the image directory)
111 in_folders=True, # save output in folders (recommended)
112 save_txt=False, # txt file for outlines in imageJ
113 save_ncolor=False,
114 ) # save ncolor version of masks for visualization and editing
115
116 def compiled_process(self, n, mask_thresh, minsize):
117 """
118 Sequentially process and save the masks in one method.
119 """
120 self.process(n, mask_thresh, minsize)
121 self.save_masks()