celldetective 1.4.2__py3-none-any.whl → 1.5.0b1__py3-none-any.whl

celldetective/tracking.py

@@ -7,8 +7,9 @@ from btrack.io.utils import localizations_to_objects
 from btrack import BayesianTracker
 
 from celldetective.measure import measure_features
-from celldetective.utils import rename_intensity_column, velocity_per_track
-from celldetective.io import interpret_tracking_configuration
+from celldetective.utils.maths import velocity_per_track
+from celldetective.utils.data_cleaning import rename_intensity_column
+from celldetective.utils.data_loaders import interpret_tracking_configuration
 
 import os
 import subprocess
@@ -95,6 +96,7 @@ def track(
 
     Examples
     --------
+
     >>> labels = np.array([[1, 1, 2, 2, 0, 0],
                                                [1, 1, 1, 2, 2, 0],
                                                [0, 0, 1, 2, 0, 0]])
@@ -473,16 +475,11 @@ def clean_trajectories(
 
     """
 
-    trajectories.reset_index
+    trajectories.reset_index(drop=True, inplace=True)
     trajectories.sort_values(
         by=[column_labels["track"], column_labels["time"]], inplace=True
     )
 
-    if minimum_tracklength > 0:
-        trajectories = filter_by_tracklength(
-            trajectories.copy(), minimum_tracklength, track_label=column_labels["track"]
-        )
-
     if np.any([remove_not_in_first, remove_not_in_last]):
         trajectories = filter_by_endpoints(
             trajectories.copy(),
@@ -503,11 +500,19 @@ def clean_trajectories(
         trajectories = interpolate_time_gaps(
             trajectories.copy(), column_labels=column_labels
         )
+        # interpolate_time_gaps might leave TRACK_ID in index of some rows or overall
+        trajectories.reset_index(drop=True, inplace=True)
 
     if interpolate_na:
         trajectories = interpolate_nan_properties(
             trajectories.copy(), track_label=column_labels["track"]
         )
+        trajectories.reset_index(drop=True, inplace=True)
+
+    if minimum_tracklength > 0:
+        trajectories = filter_by_tracklength(
+            trajectories.copy(), minimum_tracklength, track_label=column_labels["track"]
+        )
 
     trajectories = trajectories.sort_values(
         by=[column_labels["track"], column_labels["time"]]
@@ -600,10 +605,15 @@ def interpolate_nan_properties(trajectories, track_label="TRACK_ID"):
 
     """
 
-    trajectories = trajectories.groupby(track_label, group_keys=False).apply(
+    trajectories = trajectories.groupby(track_label, group_keys=True).apply(
         interpolate_per_track
     )
 
+    if track_label in trajectories.index.names:
+        trajectories = trajectories.reset_index(
+            level=0, drop=track_label in trajectories.columns
+        )
+
     return trajectories
 
 
@@ -840,11 +850,9 @@ def interpolate_time_gaps(
         lambda x: x.interpolate(method="linear")
     )
     trajectories.reset_index(drop=True, inplace=True)
-    trajectories[column_labels["time"]] = np.round(
-        (trajectories[column_labels["time"]] - pd.Timestamp("1970-01-01"))
-        / pd.Timedelta("1s"),
-        9,
-    )
+    trajectories[column_labels["time"]] = (
+        trajectories[column_labels["time"]] - pd.Timestamp("1970-01-01")
+    ).dt.total_seconds()
     # trajectories[column_labels['time']] = trajectories[column_labels['time']].astype('int64')
     trajectories.sort_values(
         by=[column_labels["track"], column_labels["time"]], inplace=True
@@ -998,7 +1006,14 @@ def extrapolate_tracks(
             [column_labels["track"], column_labels["time"]], inplace=True
         )
 
-    return trajectories
+    return trajectories[
+        [column_labels["track"], column_labels["time"]]
+        + [
+            col
+            for col in trajectories.columns
+            if col not in [column_labels["track"], column_labels["time"]]
+        ]
+    ]
 
 
 def compute_instantaneous_velocity(

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"