Improvements for Cellpose segmentation

sopa/cli/segmentation.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+from typing import Iterable
 
 import typer
 
@@ -27,6 +28,10 @@ def cellpose(
         0.2,
         help="Parameter for skimage.exposure.equalize_adapthist (applied before running cellpose)",
     ),
+    clahe_kernel_size: int = typer.Option(
+        None,
+        help="Parameter for skimage.exposure.equalize_adapthist (applied before running cellpose)",
+    ),
     gaussian_sigma: float = typer.Option(
         1, help="Parameter for scipy gaussian_filter (applied before running cellpose)"
     ),
@@ -72,6 +77,7 @@ def cellpose(
         channels,
         min_area,
         clip_limit,
+        clahe_kernel_size,
         gaussian_sigma,
         patch_index,
         patch_dir,
@@ -97,6 +103,10 @@ def generic_staining(
         0.2,
         help="Parameter for skimage.exposure.equalize_adapthist (applied before running the segmentation method)",
     ),
+    clahe_kernel_size: int = typer.Option(
+        None,
+        help="Parameter for skimage.exposure.equalize_adapthist (applied before running cellpose)",
+    ),
     gaussian_sigma: float = typer.Option(
         1,
         help="Parameter for scipy gaussian_filter (applied before running the segmentation method)",
@@ -136,6 +146,7 @@ def generic_staining(
         channels,
         min_area,
         clip_limit,
+        clahe_kernel_size,
         gaussian_sigma,
         patch_index,
         patch_dir,
@@ -149,6 +160,7 @@ def _run_staining_segmentation(
     channels: list[str],
     min_area: float,
     clip_limit: float,
+    clahe_kernel_size: int | Iterable[int] | None,
     gaussian_sigma: float,
     patch_index: int | None,
     patch_dir: str,
@@ -166,6 +178,7 @@ def _run_staining_segmentation(
         channels,
         min_area=min_area,
         clip_limit=clip_limit,
+        clahe_kernel_size=clahe_kernel_size,
         gaussian_sigma=gaussian_sigma,
     )
 

sopa/segmentation/methods.py

@@ -38,9 +38,10 @@ def cellpose_patch(
 
     cellpose_model_kwargs = cellpose_model_kwargs or {}
 
-    model = models.CellposeModel(
-        model_type=model_type, pretrained_model=pretrained_model, **cellpose_model_kwargs
-    )
+    if pretrained_model:
+        model = models.CellposeModel(pretrained_model=pretrained_model, **cellpose_model_kwargs)
+    else:
+        model = models.Cellpose(model_type=model_type, **cellpose_model_kwargs)
 
     if isinstance(channels, str) or len(channels) == 1:
         channels = [0, 0]  # gray scale

sopa/segmentation/stainings.py

@@ -2,7 +2,7 @@
 
 import logging
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Iterable
 
 import geopandas as gpd
 import numpy as np
@@ -31,6 +31,7 @@ def __init__(
         image_key: str | None = None,
         min_area: float = 0,
         clip_limit: float = 0.2,
+        clahe_kernel_size: int | Iterable[int] | None = None,
         gaussian_sigma: float = 1,
     ):
         """Generalized staining-based segmentation
@@ -71,6 +72,7 @@ def __init__(
             image_key: Optional key of `sdata` containing the image (no needed if there is only one image)
             min_area: Minimum area (in pixels^2) for a cell to be kept
             clip_limit: Parameter for skimage.exposure.equalize_adapthist (applied before running cellpose)
+            clahe_kernel_size: Parameter for skimage.exposure.equalize_adapthist (applied before running cellpose)
             gaussian_sigma: Parameter for scipy gaussian_filter (applied before running cellpose)
         """
         self.sdata = sdata
@@ -79,6 +81,7 @@ def __init__(
 
         self.min_area = min_area
         self.clip_limit = clip_limit
+        self.clahe_kernel_size = clahe_kernel_size
         self.gaussian_sigma = gaussian_sigma
 
         self.image_key, self.image = get_spatial_image(sdata, key=image_key, return_key=True)
@@ -105,8 +108,19 @@ def _run_patch(self, patch: Polygon) -> list[Polygon]:
             y=slice(bounds[1], bounds[3]),
         ).values
 
-        image = gaussian_filter(image, sigma=self.gaussian_sigma)
-        image = exposure.equalize_adapthist(image, clip_limit=self.clip_limit)
+        if self.gaussian_sigma > 0:
+            image = np.stack([gaussian_filter(c, sigma=self.gaussian_sigma) for c in image])
+        if self.clip_limit > 0:
+            image = np.stack(
+                [
+                    exposure.equalize_adapthist(
+                        c,
+                        clip_limit=self.clip_limit,
+                        kernel_size=self.clahe_kernel_size,
+                    )
+                    for c in image
+                ]
+            )
 
         if patch.area < box(*bounds).area:
             image = image * shapes.rasterize(patch, image.shape[1:], bounds)