napari-tmidas 0.2.1__py3-none-any.whl → 0.2.4__py3-none-any.whl

napari_tmidas/processing_functions/cellpose_segmentation.py

@@ -11,6 +11,7 @@ for cellular segmentation without requiring diameter parameter.
 
 Note: This requires the cellpose library to be installed.
 """
+from typing import Union
 
 import numpy as np
 
@@ -23,16 +24,15 @@ from napari_tmidas.processing_functions.cellpose_env_manager import (
 
 # Check if cellpose is directly available in this environment
 try:
-    from cellpose import core, models
+    import cellpose  # noqa: F401
 
     CELLPOSE_AVAILABLE = True
-    USE_GPU = core.use_gpu()
-    USE_DEDICATED_ENV = False
+    # Don't evaluate USE_GPU here - it should be evaluated in the cellpose environment
+
     print("Cellpose found in current environment. Using native import.")
 except ImportError:
     CELLPOSE_AVAILABLE = False
-    USE_GPU = False
-    USE_DEDICATED_ENV = True
+
     print(
         "Cellpose not found in current environment. Will use dedicated environment."
     )
@@ -88,102 +88,6 @@ def transpose_dimensions(img, dim_order):
     return img_transposed, new_dim_order, is_3d
 
 
-def run_cellpose(
-    img,
-    model,
-    channels,
-    flow_threshold=0.4,
-    cellprob_threshold=0.0,
-    dim_order="ZYX",
-    max_pixels=4000000,
-    tile_norm_blocksize=128,
-    batch_size=32,
-):
-    """
-    Run Cellpose segmentation on an image using Cellpose 4 (Cellpose-SAM).
-
-    Parameters:
-    -----------
-    img : numpy.ndarray
-        Input image
-    model : cellpose.models.Cellpose
-        Cellpose model to use
-    channels : list
-        Channels to use for segmentation [0,0] = grayscale, [1,0] = green channel, [2,0] = red channel
-    flow_threshold : float
-        Flow threshold for Cellpose
-    cellprob_threshold : float
-        Cell probability threshold
-    dim_order : str
-        Dimension order of the input image
-    max_pixels : int
-        Maximum number of pixels to process (for 2D images)
-    tile_norm_blocksize : int
-        Block size for tile normalization (new parameter in Cellpose 4)
-    batch_size : int
-        Batch size for processing multiple images or 3D slices at once
-
-    Returns:
-    --------
-    numpy.ndarray
-        Segmented image with labels
-    """
-    # First check if the image is too large (for 2D images)
-    if len(img.shape) == 2 or (len(img.shape) == 3 and "C" in dim_order):
-        # For 2D images (potentially with channels)
-        height, width = img.shape[:2]
-        total_pixels = height * width
-        if total_pixels > max_pixels:
-            raise ValueError(
-                f"Image size ({height}x{width}={total_pixels} pixels) exceeds the "
-                f"maximum size of {max_pixels} pixels. Consider downsampling."
-            )
-
-    # Transpose to expected dimension order
-    img_transposed, new_dim_order, is_3d = transpose_dimensions(img, dim_order)
-
-    # Check if we have a time series
-    has_time = "T" in new_dim_order
-
-    # Set up normalization with tile_norm_blocksize (Cellpose 4 parameter)
-    normalize = {"tile_norm_blocksize": tile_norm_blocksize}
-
-    if has_time:
-        # Handle time series - process each time point
-        n_timepoints = img_transposed.shape[0]
-        result = np.zeros(img_transposed.shape, dtype=np.uint32)
-
-        # Process each time point
-        for t in range(n_timepoints):
-            img_t = img_transposed[t]
-            masks, _, _ = model.eval(
-                img_t,
-                channels=channels,
-                flow_threshold=flow_threshold,
-                cellprob_threshold=cellprob_threshold,
-                normalize=normalize,
-                z_axis=0 if is_3d else None,
-                do_3D=is_3d,
-                batch_size=batch_size,
-            )
-            result[t] = masks
-    else:
-        # Process single time point
-        masks, _, _ = model.eval(
-            img_transposed,
-            channels=channels,
-            flow_threshold=flow_threshold,
-            cellprob_threshold=cellprob_threshold,
-            normalize=normalize,
-            z_axis=0 if is_3d else None,
-            do_3D=is_3d,
-            batch_size=batch_size,
-        )
-        result = masks
-
-    return result.astype(np.uint32)
-
-
 @BatchProcessingRegistry.register(
     name="Cellpose-SAM Segmentation",
     suffix="_labels",
@@ -208,19 +112,40 @@ def run_cellpose(
         #     "max": 3,
         #     "description": "Second channel: 0=none, 1=green, 2=red, 3=blue",
         # },
+        "diameter": {
+            "type": float,
+            "default": 0.0,
+            "min": 0.0,
+            "max": 200.0,
+            "description": "Optional. Only required if your ROI diameter is outside the range 7.5–120. Set to 0 to leave unset (recommended for most users). Cellpose-SAM is trained for diameters in this range.",
+        },
         "flow_threshold": {
             "type": float,
             "default": 0.4,
             "min": 0.1,
             "max": 0.9,
-            "description": "Flow threshold for Cellpose segmentation",
+            "description": "Flow threshold (cell detection sensitivity)",
         },
         "cellprob_threshold": {
             "type": float,
             "default": 0.0,
             "min": -6.0,
             "max": 6.0,
-            "description": "Cell probability threshold (Cellpose 4 parameter)",
+            "description": "Cell probability threshold (increase to reduce splits)",
+        },
+        "anisotropy": {
+            "type": float,
+            "default": 1.0,
+            "min": 0.1,
+            "max": 10.0,
+            "description": "Optional rescaling factor (3D; Z step[um] / X pixel res [um])",
+        },
+        "flow3D_smooth": {
+            "type": int,
+            "default": 0,
+            "min": 0,
+            "max": 10,
+            "description": "Smooth flow with gaussian filter (stdev)",
         },
         "tile_norm_blocksize": {
             "type": int,
@@ -236,11 +161,6 @@ def run_cellpose(
             "max": 128,
             "description": "Batch size for processing multiple images/slices at once",
         },
-        "force_dedicated_env": {
-            "type": bool,
-            "default": False,
-            "description": "Force using dedicated environment even if Cellpose is available",
-        },
     },
 )
 def cellpose_segmentation(
@@ -250,9 +170,12 @@ def cellpose_segmentation(
     channel_2: int = 0,
     flow_threshold: float = 0.4,
     cellprob_threshold: float = 0.0,
+    anisotropy: Union[float, None] = None,
+    flow3D_smooth: int = 0,
     tile_norm_blocksize: int = 128,
     batch_size: int = 32,
-    force_dedicated_env: bool = False,
+    diameter: float = 0.0,
+    _source_filepath: str = None,  # Hidden parameter for zarr optimization
 ) -> np.ndarray:
     """
     Run Cellpose 4 (Cellpose-SAM) segmentation on an image.
@@ -282,110 +205,128 @@ def cellpose_segmentation(
         Block size for tile normalization (Cellpose 4 parameter) (default: 128)
     batch_size : int
         Batch size for processing multiple images/slices at once (default: 32)
-    force_dedicated_env : bool
-        Force using dedicated environment even if Cellpose is available (default: False)
 
     Returns:
     --------
     numpy.ndarray
         Segmented image with instance labels
     """
+    # Diameter parameter guidance:
+    # Cellpose-SAM is trained for ROI diameters 7.5–120. Only set diameter if your images have objects outside this range (e.g., diameter <7.5 or >120). Otherwise, leave as None.
+    # Cellpose 4 handles normalization internally via percentile-based normalization
+    # It accepts uint8, uint16, float32, float64 - no pre-conversion needed!
+    # The normalize=True parameter (default) will convert to float and normalize
+    # to 1st-99th percentile range internally
+
+    # Handle TZYX data by processing each timepoint separately
+    if "T" in dim_order and image.ndim == 4:
+        print(
+            f"Detected TZYX data with shape {image.shape}. Processing each timepoint separately..."
+        )
+        t_axis = dim_order.index("T")
+        num_timepoints = image.shape[t_axis]
+
+        # Move T axis to front if not already there
+        if t_axis != 0:
+            axes_order = list(range(image.ndim))
+            axes_order.insert(0, axes_order.pop(t_axis))
+            image = np.transpose(image, axes_order)
+
+        # Process each timepoint
+        results = []
+        for t in range(num_timepoints):
+            print(f"Processing timepoint {t+1}/{num_timepoints}...")
+            timepoint_image = image[t]
+            # Remove 'T' from dim_order for single timepoint
+            timepoint_dim_order = dim_order.replace("T", "")
+
+            # Recursively call this function for the single timepoint
+            timepoint_result = cellpose_segmentation(
+                timepoint_image,
+                dim_order=timepoint_dim_order,
+                channel_1=channel_1,
+                channel_2=channel_2,
+                flow_threshold=flow_threshold,
+                cellprob_threshold=cellprob_threshold,
+                anisotropy=anisotropy,
+                flow3D_smooth=flow3D_smooth,
+                tile_norm_blocksize=tile_norm_blocksize,
+                batch_size=batch_size,
+                _source_filepath=None,  # Don't use zarr optimization for timepoints
+            )
+            results.append(timepoint_result)
+
+        # Stack results back together
+        result = np.stack(results, axis=0)
+        print(f"Completed processing all {num_timepoints} timepoints.")
+        return result
+
     # Convert channel parameters to Cellpose channels list
     # channels = [channel_1, channel_2]
     channels = [0, 0]  # limit script to single channel
-    # Determine whether to use dedicated environment
-    use_env = force_dedicated_env or USE_DEDICATED_ENV
 
-    if use_env:
-        print("Using dedicated Cellpose environment...")
+    # First check if the environment exists, create if not
+    if not is_env_created():
+        print(
+            "Creating dedicated Cellpose environment (this may take a few minutes)..."
+        )
+        create_cellpose_env()
+        print("Environment created successfully.")
 
-        # First check if the environment exists, create if not
-        if not is_env_created():
-            print(
-                "Creating dedicated Cellpose environment (this may take a few minutes)..."
-            )
-            create_cellpose_env()
-            print("Environment created successfully.")
+    # Check if we can use zarr optimization
+    use_zarr_direct = _source_filepath and _source_filepath.lower().endswith(
+        ".zarr"
+    )
 
-        # Prepare arguments for the Cellpose function
+    if use_zarr_direct:
+        print(f"Using optimized zarr processing for: {_source_filepath}")
+        # Prepare arguments for direct zarr processing
+        is_3d = "Z" in dim_order
         args = {
-            "image": image,
+            "zarr_path": _source_filepath,
+            "zarr_key": None,  # Could be enhanced to support specific keys
             "channels": channels,
             "flow_threshold": flow_threshold,
             "cellprob_threshold": cellprob_threshold,
+            "flow3D_smooth": flow3D_smooth,
+            "anisotropy": anisotropy,
             "normalize": {"tile_norm_blocksize": tile_norm_blocksize},
             "batch_size": batch_size,
-            "use_gpu": USE_GPU,
-            "do_3D": "Z" in dim_order,
-            "z_axis": 0 if "Z" in dim_order else None,
+            "diameter": diameter,
+            "use_gpu": True,  # Let cellpose environment detect GPU
+            "do_3D": is_3d,
+            "z_axis": 0 if is_3d else None,
+            "channel_axis": None,  # No channel axis for single-channel grayscale images
         }
-
-        # Run Cellpose in the dedicated environment
-        print("Running Cellpose model in dedicated environment...")
-        result = run_cellpose_in_env("eval", args)
-        print(f"Segmentation complete. Found {np.max(result)} objects.")
-        return result
-
     else:
-        print("Running Cellpose model in current environment...")
-        # Initialize Cellpose model in current environment
-        model = models.CellposeModel(gpu=USE_GPU)
-
-    # Print status information
-    gpu_status = "GPU" if USE_GPU else "CPU"
-    print(f"Using Cellpose on {gpu_status}")
-    print(
-        f"Processing image with shape {image.shape}, dimension order: {dim_order}"
-    )
-
-    # Run segmentation
-    try:
-        result = run_cellpose(
-            image,
-            model,
-            channels,
-            flow_threshold,
-            cellprob_threshold,
-            dim_order,
-            tile_norm_blocksize=tile_norm_blocksize,
-            batch_size=batch_size,
-        )
+        # Prepare arguments for the Cellpose function (legacy numpy array mode)
+        is_3d = "Z" in dim_order
+        args = {
+            "image": image,
+            "channels": channels,
+            "flow_threshold": flow_threshold,
+            "cellprob_threshold": cellprob_threshold,
+            "flow3D_smooth": flow3D_smooth,
+            "anisotropy": anisotropy,
+            "normalize": {"tile_norm_blocksize": tile_norm_blocksize},
+            "batch_size": batch_size,
+            "diameter": diameter,
+            "use_gpu": True,  # Let cellpose environment detect GPU
+            "do_3D": is_3d,
+            "z_axis": 0 if is_3d else None,
+            "channel_axis": None,  # No channel axis for single-channel grayscale images
+        }
+    # Run Cellpose in the dedicated environment
+    print("Running Cellpose model in dedicated environment...")
+    result = run_cellpose_in_env("eval", args)
+    print(f"Segmentation complete. Found {np.max(result)} objects.")
 
-        print(f"Segmentation complete. Found {np.max(result)} objects.")
-        return result
+    # Ensure result is uint32 for proper label detection in napari
+    # Cellpose typically returns int32, but uint32 is preferred for labels
+    if result.dtype != np.uint32:
+        result = result.astype(np.uint32)
 
-    except (Exception, MemoryError) as e:
-        print(f"Error during segmentation in current environment: {str(e)}")
-
-        # If we haven't already tried using the dedicated environment, try that as a fallback
-        if not USE_DEDICATED_ENV and not force_dedicated_env:
-            print("Attempting fallback to dedicated Cellpose environment...")
-            try:
-                args = {
-                    "image": image,
-                    "channels": channels,
-                    "flow_threshold": flow_threshold,
-                    "cellprob_threshold": cellprob_threshold,
-                    "normalize": {"tile_norm_blocksize": tile_norm_blocksize},
-                    "batch_size": batch_size,
-                    "use_gpu": USE_GPU,
-                    "do_3D": "Z" in dim_order,
-                    "z_axis": 0 if "Z" in dim_order else None,
-                }
-
-                if not is_env_created():
-                    create_cellpose_env()
-
-                result = run_cellpose_in_env("eval", args)
-                print(f"Fallback successful. Found {np.max(result)} objects.")
-                return result
-            except (Exception, MemoryError) as fallback_e:
-                print(f"Fallback also failed: {str(fallback_e)}")
-                raise Exception(
-                    f"Both direct and dedicated environment approaches failed: {str(e)} | {str(fallback_e)}"
-                ) from fallback_e
-        else:
-            raise
+    return result
 
 
 # Update cellpose_env_manager.py to install Cellpose 4