  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()