Adding segmentation utility functions

viscy/analysis/segmentation.py

@@ -0,0 +1,118 @@
+import torch
+import numpy as np
+import click
+from cellpose import models
+from skimage.exposure import rescale_intensity, equalize_adapthist
+from skimage.util import invert
+from numpy.typing import ArrayLike
+
+
+def nuc_mem_segmentation_cellposemodel_3D(
+    czyx_data: ArrayLike, zyx_slicing: tuple[slice, slice, slice], **cellpose_kwargs
+):
+    """
+    Segment nuclei and membranes using Cellpose 3D model.
+
+    """
+
+    Z_slice = zyx_slicing[0]
+    Y_slice = zyx_slicing[1]
+    X_slice = zyx_slicing[2]
+    czyx_data = czyx_data[:, Z_slice, Y_slice, X_slice]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    segmentation_stack = np.zeros_like(czyx_data)
+    click.echo(f"Segmentation Stack shape {segmentation_stack.shape}")
+    cellpose_params = cellpose_kwargs["cellpose_kwargs"]
+    c_idx = 0
+    if "nucleus_kwargs" in cellpose_params:
+        click.echo("Segmenting Nuclei")
+        nuc_seg_kwargs = cellpose_params["nucleus_kwargs"]
+
+        model_nucleus_3D = models.CellposeModel(
+            model_type=cellpose_params["nuc_model_path"],
+            # net_avg=True, #Note removed CP3.0
+            gpu=True,
+            device=torch.device(device),
+        )
+        nuc_segmentation, _, _ = model_nucleus_3D.eval(czyx_data, **nuc_seg_kwargs)
+        segmentation_stack[c_idx] = nuc_segmentation.astype(np.uint16)
+        c_idx += 1
+    if "membrane_kwargs" in cellpose_params:
+        click.echo("Segmenting Membrane")
+        mem_seg_kwargs = cellpose_params["membrane_kwargs"]
+
+        model_membrane_3D = models.CellposeModel(
+            model_type=cellpose_params["mem_model_path"],
+            # net_avg=True,
+            gpu=True,
+            device=torch.device(device),
+        )
+        c_idx_mem, c_idx_nuc = mem_seg_kwargs["channels"]
+        mem_segmentation, _, _ = model_membrane_3D.eval(czyx_data, **mem_seg_kwargs)
+        segmentation_stack[c_idx] = mem_segmentation.astype(np.uint16)
+
+    return segmentation_stack
+
+
+def nuc_mem_cp_segmentation_clahe_3D(
+    czyx_data: ArrayLike, zyx_slicing: tuple, clahe_kwargs, **cellpose_kwargs
+):
+"""
+    Segment nuclei and membranes using Cellpose 3D model with CLAHE applied to the input data.
+"""
+
+    Z_slice = zyx_slicing[0]
+    Y_slice = zyx_slicing[1]
+    X_slice = zyx_slicing[2]
+    czyx_data = czyx_data[:, Z_slice, Y_slice, X_slice]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    segmentation_stack = np.zeros_like(czyx_data, dtype=np.uint16)
+    click.echo(f"Segmentation Stack shape {segmentation_stack.shape}")
+    cellpose_params = cellpose_kwargs["cellpose_kwargs"]
+    # clahe_kwargs = clahe_kwargs['clahe']
+    c_idx = 0
+    if "nucleus_kwargs" in cellpose_params:
+        click.echo("Segmenting Nuclei")
+        nuc_seg_kwargs = cellpose_params["nucleus_kwargs"]
+
+        model_nucleus_3D = models.CellposeModel(
+            model_type=cellpose_params["nuc_model_path"],
+            # net_avg=True, #Note removed CP3.0
+            gpu=True,
+            device=torch.device(device),
+        )
+        # Apply CLAHE before cellpose
+        if "clahe_nuc" in clahe_kwargs:
+            click.echo("Applying CLAHE to Nuclei")
+            nuc_clahe = clahe_kwargs["clahe_nuc"]
+            czyx_data[c_idx] = rescale_intensity(czyx_data[c_idx], out_range=(0.0, 1.0))
+            czyx_data[c_idx] = equalize_adapthist(czyx_data[c_idx], **nuc_clahe)
+        nuc_segmentation, _, _ = model_nucleus_3D.eval(czyx_data, **nuc_seg_kwargs)
+        segmentation_stack[c_idx] = nuc_segmentation.astype(np.uint16)
+        c_idx += 1
+    if "membrane_kwargs" in cellpose_params:
+        click.echo("Segmenting Membrane")
+        mem_seg_kwargs = cellpose_params["membrane_kwargs"]
+
+        if "clahe_mem" in clahe_kwargs:
+            click.echo("Applying CLAHE to Membrane")
+            mem_clahe = clahe_kwargs["clahe_mem"]
+            czyx_data[c_idx] = rescale_intensity(
+                invert(czyx_data[c_idx]), out_range=(0.0, 1.0)
+            )
+            czyx_data[c_idx] = equalize_adapthist(czyx_data[c_idx], **mem_clahe)
+        model_membrane_3D = models.CellposeModel(
+            model_type=cellpose_params["mem_model_path"],
+            # net_avg=True,
+            gpu=True,
+            device=torch.device(device),
+        )
+        c_idx_mem, c_idx_nuc = mem_seg_kwargs["channels"]
+        mem_segmentation, _, _ = model_membrane_3D.eval(czyx_data, **mem_seg_kwargs)
+        segmentation_stack[c_idx] = mem_segmentation.astype(np.uint16)
+
+    return segmentation_stack

viscy/analysis/settings/segmentation.yml

@@ -0,0 +1,24 @@
+mem_model_path: "/hpc/projects/jacobo_group/Code/timelapse_seg_tracking_pipeline/3_segmentation/membrane/cellpose_2Chan_scratch_2024_04_30_11_12_00"
+membrane_kwargs:
+  diameter: 65
+  channels:
+    - 2
+    - 1
+  cellprob_threshold: 0.4
+  invert: false
+  do_3D: true
+  anisotropy: 3.26
+  min_size: 8000
+
+nuc_model_path: "/hpc/projects/jacobo_group/projects/cellpose/Nuclei/Deconvolved/Fine_Tune/models/cellpose_Slices_decon_nuclei_nuclei_v7_2023_06_28_16_54"
+nucleus_kwargs:
+  diameter: 60
+  channels:
+    - 1
+    - 0
+  cellprob_threshold: 0.0
+  invert: false
+  do_3D: true
+  anisotropy: 3.26
+  min_size: 8000
+