Skip to content

Segmentation Utilities#

This is a collection of wrapper functions for omnipose.

Run the specified omnipose model on the input image(s)

Parameters:

Name Type Description Default
images np.array

The input image(s) to segment. If a single image is provided, it should be a 2D array. If multiple images are provided, they should be a 3D array with shape (n_images, height, width).

 required
model_name str

The name of the model to use.

 required
gpu_option bool

Whether to use the GPU for segmentation.

 required
chans tuple or None

The channels to use for segmentation. If None, the first channel will be used.

 None
filter_options dict

A dictionary containing the filter options for the segmentation. Related to the omnipose package. See this link for more info: https://omnipose.readthedocs.io/examples/mono_channel_bact.html#run-segmentation

 {}

Returns:

Name Type Description
masks np.ndarray

The segmented masks. If a single image is provided, it will be a 2D array. If multiple images are provided, it will be a 3D array with shape (n_images, height, width).
flows np.ndarray

The flow outputs. If a single image is provided, it will be a 2D array. If multiple images are provided, it will be a 3D array with shape (n_images, height, width).
Source code in src/micromorph/segmentation/segmentation_tools.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def run_omnipose(images: np.array, model_name: str, gpu_option: bool, chans: tuple or None =None, filter_options:
dict ={}) -> Union[np.array, np.array]:
    """
    Run the specified omnipose model on the input image(s)

    Parameters
    ----------
    images: np.ndarray
        The input image(s) to segment. If a single image is provided, it should be a 2D array. If multiple images are provided, they should be a 3D array with shape (n_images, height, width).
    model_name: str
        The name of the model to use.
    gpu_option: bool
        Whether to use the GPU for segmentation.
    chans: tuple or None
        The channels to use for segmentation. If None, the first channel will be used.
    filter_options: dict
        A dictionary containing the filter options for the segmentation. Related to the omnipose package. See this 
        link for more info: https://omnipose.readthedocs.io/examples/mono_channel_bact.html#run-segmentation

    Returns
    -------
    masks: np.ndarray
        The segmented masks. If a single image is provided, it will be a 2D array. If multiple images are provided, it will be a 3D array with shape (n_images, height, width).
    flows: np.ndarray
        The flow outputs. If a single image is provided, it will be a 2D array. If multiple images are provided, it will be a 3D array with shape (n_images, height, width).
    """
    # Should also add to this function the possibility to:
    model = models.CellposeModel(gpu=gpu_option, model_type=model_name)

    if chans is None:
        chans = [0, 0]  # If user didn't specify a channel, assume you need to use the first channel

    # TODO: params should be a dictionary which you can edit

    # define parameters - this is to become an input parameter
    params = {'channels': chans,  # always define this with the model
              'rescale': None,  # upscale or downscale your images, None = no rescaling
              'mask_threshold': 1,  # erode or dilate masks with higher or lower values -1
              'flow_threshold': 0.2,
              # default is .4, but only needed if there are spurious masks to clean up; slows down output
              'transparency': True,  # transparency in flow output
              'omni': True,  # we can turn off Omnipose mask reconstruction, not advised
              'cluster': True,  # use DBSCAN clustering
              'resample': True,  # whether to run dynamics on rescaled grid or original grid
              # 'verbose': False, # turn on if you want to see more output
              'tile': False,  # average the outputs from flipped (augmented) images; slower, usually not needed
              'niter': None,
              # None lets Omnipose calculate # of Euler iterations (usually <20) but you can tune it for over/under segmentation
              'augment': False,  # Can optionally rotate the image and average outputs, usually not needed
              'affinity_seg': False,  # new feature, stay tuned...
              }

    tic = time.time()
    if len(images.shape) == 2:
        try:
            masks, flows, styles = model.eval(images, **params)
        except ValueError as e:
            logging.error(f"Error in Omnipose segmentation: {e}")
            logging.info("Error in Omnipose segmentation. A blank mask will be returned.")
            masks = np.zeros_like(images)
            flows = np.zeros_like(images)
    else:
        n_images = len(images)
        n = range(n_images)

        masks_list, flows_list, styles = model.eval([images[i] for i in n], **params)

        # Convert list of masks to 3D array
        masks = np.dstack(masks_list)  # Transpose to match the original shape of the stack
        masks = np.moveaxis(masks, -1, 0)

        flows = np.copy(masks)

        for i, sublist in enumerate(flows_list):
            flows[i] = sublist[-1]  # -1

    net_time = time.time() - tic
    logging.info('Total segmentation time: {}s'.format(net_time))

    if filter_options:
        masks = filter_mask(masks, filter_options)

    return masks, flows

Function to filter the segmentation mask basaed on basic morphological criteria such as area, length and width.

Parameters:

Name Type Description Default
mask np.array

The mask to filter. If a single mask is provided, it should be a 2D array. If multiple masks are provided, they should be a 3D array with shape (n_masks, height, width).

 required
options dict

A dictionary containing the filter options for the segmentation. The options are: - min_area: The minimum area of the bacteria. - max_area: The maximum area of the bacteria. - min_length: The minimum length of the bacteria. - max_length: The maximum length of the bacteria. - min_width: The minimum width of the bacteria. - max_width: The maximum width of the bacteria.

 {}

Returns:

Name Type Description
filtered_mask np.ndarray

The filtered mask. If a single mask is provided, it will be a 2D array. If multiple masks are provided, it will be a 3D array with shape (n_masks, height, width).
Source code in src/micromorph/segmentation/segmentation_tools.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def filter_mask(mask:np.array, options: dict ={}) -> np.array:
    """
    Function to filter the segmentation mask basaed on basic morphological criteria such as area, length and width.

    Parameters
    ----------
    mask: np.ndarray
        The mask to filter. If a single mask is provided, it should be a 2D array. If multiple masks are provided, they should be a 3D array with shape (n_masks, height, width).
    options: dict
        A dictionary containing the filter options for the segmentation. The options are:
        - min_area: The minimum area of the bacteria.
        - max_area: The maximum area of the bacteria.
        - min_length: The minimum length of the bacteria.
        - max_length: The maximum length of the bacteria.
        - min_width: The minimum width of the bacteria.
        - max_width: The maximum width of the bacteria.

    Returns
    -------
    filtered_mask: np.ndarray
        The filtered mask. If a single mask is provided, it will be a 2D array. If multiple masks are provided, it will be a 3D array with shape (n_masks, height, width).
    """
    # Based on this: https://forum.image.sc/t/need-help-with-regionprops/30107/2

    # Get values from dictionary if they exist, otherwise use default values
    min_area = options.get("min_area", None)
    max_area = options.get("max_area", None)

    min_length = options.get("min_length", None)
    max_length = options.get("max_length", None)

    min_width = options.get("min_width", None)
    max_width = options.get("max_width", None)

    # Not sure if needed ---
    # label_mask = label(mask)
    label_mask = np.copy(mask)

    if len(label_mask.shape) == 2:
        # Start by getting regionprops
        regions = regionprops(label_mask)

        # Create a list to store the indices of the bacteria that meet the criteria
        bacteria_indices = []
        # Create a new mask with only the bacteria that meet the criteria
        filtered_mask = np.zeros_like(mask)

        # Loop through the regions
        for i, region in enumerate(regions):
            # Get the area, length and width of the region
            area = region.area
            length = region.axis_major_length
            width = region.axis_minor_length

            # Check if the region meets the criteria
            if (min_area is None or area >= min_area) and (max_area is None or area <= max_area) and \
                    (min_length is None or length >= min_length) and (max_length is None or length <= max_length) and \
                    (min_width is None or width >= min_width) and (max_width is None or width <= max_width):
                # If the region meets the criteria, add the index to the list
                bacteria_indices.append(region.label)

        for i in bacteria_indices:
            filtered_mask[label_mask == i] = i
    elif len(label_mask.shape) == 3:
        # Create stack to store the filtered masks
        filtered_mask = np.zeros_like(mask)

        # Loop through the frames
        for i, current_mask in enumerate(label_mask):
            current_filtered_mask = filter_mask(current_mask, options=options)

            filtered_mask[i] = current_filtered_mask
    else:
        raise ValueError("Mask must be 2D or 3D.")

    return filtered_mask