celldetective 1.4.1.post1__py3-none-any.whl → 1.5.0b0__py3-none-any.whl

celldetective/utils/cellpose_utils/__init__.py

@@ -0,0 +1,133 @@
+from pathlib import Path
+from typing import Union, Optional
+
+import numpy as np
+
+
+def _segment_image_with_cellpose_model(
+    img,
+    model=None,
+    diameter=None,
+    cellprob_threshold=None,
+    flow_threshold=None,
+    channel_axis=-1,
+):
+    """
+    Segments an input image using a Cellpose model.
+
+    This function applies a preloaded Cellpose model to segment an input image and returns the resulting labeled mask.
+    The image is rearranged into the format expected by the Cellpose model, with the specified channel axis moved to the first dimension.
+
+    Parameters
+    ----------
+    img : ndarray
+            The input image to be segmented. It is expected to have a channel axis specified by `channel_axis`.
+    model : CellposeModel, optional
+            A preloaded Cellpose model instance used for segmentation.
+    diameter : float, optional
+            The diameter of objects to segment. If `None`, the model's default diameter is used.
+    cellprob_threshold : float, optional
+            The threshold for the probability of cells used during segmentation. If `None`, the default threshold is used.
+    flow_threshold : float, optional
+            The threshold for flow error during segmentation. If `None`, the default threshold is used.
+    channel_axis : int, optional
+            The axis of the input image that represents the channels. Default is `-1` (channel-last format).
+
+    Returns
+    -------
+    ndarray
+            A labeled mask of the same spatial dimensions as the input image, with segmented regions assigned unique
+            integer labels. The dtype of the mask is `uint16`.
+
+    Notes
+    -----
+    - The `img` array is internally rearranged to move the specified `channel_axis` to the first dimension to comply
+      with the Cellpose model's input requirements.
+    - Ensure the provided `model` is a properly initialized Cellpose model instance.
+    - Parameters `diameter`, `cellprob_threshold`, and `flow_threshold` allow fine-tuning of the segmentation process.
+    """
+
+    img = np.moveaxis(img, channel_axis, 0)
+    lbl, _, _ = model.eval(
+        img,
+        diameter=diameter,
+        cellprob_threshold=cellprob_threshold,
+        flow_threshold=flow_threshold,
+        channels=None,
+        normalize=False,
+        model_loaded=True,
+    )
+
+    return lbl.astype(np.uint16)
+
+
+def _prep_cellpose_model(
+    model_name: str,
+    path: Union[str, Path],
+    use_gpu: bool = False,
+    n_channels: int = 2,
+    scale: Optional[float] = None,
+):
+    """
+    Prepares and loads a Cellpose model for segmentation tasks.
+
+    This function initializes a Cellpose model with the specified parameters, configures GPU usage if available,
+    and calculates or applies a scaling factor for the model based on image resolution.
+
+    Parameters
+    ----------
+    model_name : str
+            The name of the pretrained Cellpose model to load.
+    path : str
+            The directory where the model is stored.
+    use_gpu : bool, optional
+            If `True`, the model will use GPU acceleration for computations. Default is `False`.
+    n_channels : int, optional
+            The number of input channels expected by the model. Default is `2`.
+    scale : float, optional
+            A scaling factor to adjust the model's output to match the image resolution. If not provided, the scale is
+            automatically calculated based on the model's diameter parameters.
+
+    Returns
+    -------
+    tuple
+            - model : CellposeModel
+                    The loaded Cellpose model configured with the specified parameters.
+            - scale_model : float
+                    The scaling factor applied to the model, calculated or provided.
+
+    Notes
+    -----
+    - Ensure the Cellpose package is installed and the model files are correctly stored in the provided path.
+    - GPU support depends on the availability of compatible hardware and software setup.
+    - The scale is calculated as `(diam_mean / diam_labels)` if `scale` is not provided, where `diam_mean` and
+      `diam_labels` are attributes of the model.
+    """
+
+    import torch
+
+    if not use_gpu:
+        device = torch.device("cpu")
+    else:
+        device = torch.device("cuda")
+
+    from cellpose.models import CellposeModel
+
+    model = CellposeModel(
+        gpu=use_gpu,
+        device=device,
+        pretrained_model=path + model_name,
+        model_type=None,
+        nchan=n_channels,
+    )  # diam_mean=30.0,
+    if scale is None:
+        scale_model = model.diam_mean / model.diam_labels
+    else:
+        scale_model = scale * model.diam_mean / model.diam_labels
+
+    print(f"Cell size in model: {model.diam_mean} pixels...")
+    print(f"Cell size in training set: {model.diam_labels} pixels...")
+    print(f"Rescaling factor to apply: {scale_model}...")
+
+    print(f"Cellpose model {model_name} successfully loaded...")
+    return model, scale_model

celldetective/utils/color_mappings.py

@@ -0,0 +1,42 @@
+def color_from_status(status, recently_modified=False):
+
+    if not recently_modified:
+        if status == 0:
+            return "tab:blue"
+        elif status == 1:
+            return "tab:red"
+        elif status == 2:
+            return "yellow"
+        else:
+            return "k"
+    else:
+        if status == 0:
+            return "tab:cyan"
+        elif status == 1:
+            return "tab:orange"
+        elif status == 2:
+            return "tab:olive"
+        else:
+            return "k"
+
+
+def color_from_class(cclass, recently_modified=False):
+
+    if not recently_modified:
+        if cclass == 0:
+            return "tab:red"
+        elif cclass == 1:
+            return "tab:blue"
+        elif cclass == 2:
+            return "yellow"
+        else:
+            return "k"
+    else:
+        if cclass == 0:
+            return "tab:orange"
+        elif cclass == 1:
+            return "tab:cyan"
+        elif cclass == 2:
+            return "tab:olive"
+        else:
+            return "k"