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

napari_tmidas/processing_functions/cellpose_env_manager.py

@@ -1,6 +1,7 @@
 # processing_functions/cellpose_env_manager.py
 """
 This module manages a dedicated virtual environment for Cellpose.
+Updated to support Cellpose 4 (Cellpose-SAM) installation.
 """
 
 import os
@@ -59,18 +60,47 @@ def create_cellpose_env():
     # Path to the Python executable in the new environment
     env_python = get_env_python_path()
 
-    # Install numpy first to ensure correct version
-    print("Installing NumPy...")
-    subprocess.check_call(
-        [env_python, "-m", "pip", "install", "numpy>=1.24,<1.25"]
+    # # Install numpy first to ensure correct version
+    # print("Installing NumPy...")
+    # subprocess.check_call(
+    #     [env_python, "-m", "pip", "install", "--upgrade", "pip"]
+    # )
+    # subprocess.check_call(
+    #     [env_python, "-m", "pip", "install", "numpy>=1.24,<1.25"]
+    # )
+
+    # Install cellpose 4 and other dependencies
+    print(
+        "Installing Cellpose 4 (Cellpose-SAM) in the dedicated environment..."
     )
-
-    # Install cellpose and other dependencies
-    print("Installing Cellpose in the dedicated environment...")
     subprocess.check_call([env_python, "-m", "pip", "install", "cellpose"])
 
-    # Install tifffile for image handling
-    subprocess.check_call([env_python, "-m", "pip", "install", "tifffile"])
+    # # Install PyTorch - needed for GPU acceleration
+    # print("Installing PyTorch for GPU acceleration...")
+    # subprocess.check_call([
+    #     env_python, "-m", "pip", "install",
+    #     "torch", "torchvision"
+    # ])
+
+    # # Install tifffile for image handling
+    # subprocess.check_call([env_python, "-m", "pip", "install", "tifffile"])
+
+    # Check if installation was successful
+    try:
+        # Run a command to check if cellpose can be imported and GPU is available
+        result = subprocess.run(
+            [
+                env_python,
+                "-c",
+                "from cellpose import core; print(f'GPU available: {core.use_gpu()}')",
+            ],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        print(f"Cellpose environment check: {result.stdout.strip()}")
+    except subprocess.CalledProcessError as e:
+        print(f"Warning: Cellpose installation check failed: {e.stderr}")
 
     print("Cellpose environment created successfully.")
     return env_python
@@ -109,26 +139,32 @@ def run_cellpose_in_env(func_name, args_dict):
         tifffile.imwrite(input_file.name, args_dict["image"])
 
         # Prepare a temporary script to run Cellpose
+        # Updated to use Cellpose 4 parameters
         script = f"""
 import numpy as np
-from cellpose import models
+from cellpose import models, core
 import tifffile
 
+
+
 # Load image
 image = tifffile.imread('{input_file.name}')
 
 # Create and run model
-model = models.Cellpose(
-    gpu={args_dict.get('use_gpu', True)},
-    model_type='{args_dict.get('model_type', 'cyto3')}'
-)
+model = models.CellposeModel(
+    gpu={args_dict.get('use_gpu', True)})
+
+# Prepare normalization parameters (Cellpose 4)
+normalize = {args_dict.get('normalize', {'tile_norm_blocksize': 128})}
 
-# Perform segmentation
-masks, *_ = model.eval(
+# Perform segmentation with Cellpose 4 parameters
+masks, flows, styles = model.eval(
     image,
-    diameter={args_dict.get('diameter', 30.0)},
-    flow_threshold={args_dict.get('flow_threshold', 0.4)},
     channels={args_dict.get('channels', [0, 0])},
+    flow_threshold={args_dict.get('flow_threshold', 0.4)},
+    cellprob_threshold={args_dict.get('cellprob_threshold', 0.0)},
+    batch_size={args_dict.get('batch_size', 32)},
+    normalize=normalize,
     do_3D={args_dict.get('do_3D', False)},
     z_axis={args_dict.get('z_axis', 0)} if {args_dict.get('do_3D', False)} else None
 )
@@ -147,10 +183,9 @@ tifffile.imwrite('{output_file.name}', masks)
         result = subprocess.run(
             [env_python, script_file.name], capture_output=True, text=True
         )
-
+        print("Stdout:", result.stdout)
         # Check for errors
         if result.returncode != 0:
-            print("Stdout:", result.stdout)
             print("Stderr:", result.stderr)
             raise RuntimeError(
                 f"Cellpose segmentation failed: {result.stderr}"

napari_tmidas/processing_functions/cellpose_segmentation.py

@@ -6,10 +6,11 @@ This module provides functionality to automatically segment cells or nuclei in i
 using the Cellpose deep learning-based segmentation toolkit. It supports both 2D and 3D images,
 various dimension orders, and handles time series data.
 
+Updated to support Cellpose 4 (Cellpose-SAM) which offers improved generalization
+for cellular segmentation without requiring diameter parameter.
+
 Note: This requires the cellpose library to be installed.
 """
-import os
-import sys
 
 import numpy as np
 
@@ -91,13 +92,15 @@ def run_cellpose(
     img,
     model,
     channels,
-    diameter,
     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.
+    Run Cellpose segmentation on an image using Cellpose 4 (Cellpose-SAM).
 
     Parameters:
     -----------
@@ -107,14 +110,18 @@ def run_cellpose(
         Cellpose model to use
     channels : list
         Channels to use for segmentation [0,0] = grayscale, [1,0] = green channel, [2,0] = red channel
-    diameter : float
-        Diameter of objects to segment
     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:
     --------
@@ -138,6 +145,9 @@ def run_cellpose(
     # 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]
@@ -146,67 +156,58 @@ def run_cellpose(
         # Process each time point
         for t in range(n_timepoints):
             img_t = img_transposed[t]
-            mask, _, _, _ = model.eval(
+            masks, _, _ = model.eval(
                 img_t,
-                diameter=diameter,
-                flow_threshold=flow_threshold,
                 channels=channels,
+                flow_threshold=flow_threshold,
+                cellprob_threshold=cellprob_threshold,
+                normalize=normalize,
                 z_axis=0 if is_3d else None,
                 do_3D=is_3d,
-                niter=2000,  # Maximum iterations
+                batch_size=batch_size,
             )
-            result[t] = mask
+            result[t] = masks
     else:
         # Process single time point
-        result, _, _, _ = model.eval(
+        masks, _, _ = model.eval(
             img_transposed,
-            diameter=diameter,
-            flow_threshold=flow_threshold,
             channels=channels,
+            flow_threshold=flow_threshold,
+            cellprob_threshold=cellprob_threshold,
+            normalize=normalize,
             z_axis=0 if is_3d else None,
             do_3D=is_3d,
-            niter=2000,  # Maximum iterations
+            batch_size=batch_size,
         )
+        result = masks
 
     return result.astype(np.uint32)
 
 
 @BatchProcessingRegistry.register(
-    name="Segment cells or nuclei (Cellpose)",
+    name="Cellpose-SAM Segmentation",
     suffix="_labels",
-    description="Automatic instance segmentation using Cellpose 3.0 (dedicated environment)",
+    description="Automatic instance segmentation using Cellpose 4 (Cellpose-SAM) with improved generalization.",
     parameters={
-        "model_type": {
-            "type": str,
-            "default": "cyto3",
-            "description": "Cellpose model type: 'cyto'/'cyto2'/'cyto3' for cells, 'nuclei' for nuclei",
-        },
-        "diameter": {
-            "type": float,
-            "default": 40.0,
-            "min": 5.0,
-            "max": 1000.0,
-            "description": "Approximate diameter of objects to segment (pixels)",
-        },
         "dim_order": {
             "type": str,
             "default": "YX",
             "description": "Dimension order of the input (e.g., 'YX', 'ZYX', 'TZYX')",
         },
-        "channel_1": {
-            "type": int,
-            "default": 0,
-            "min": 0,
-            "max": 3,
-            "description": "First channel: 0=grayscale, 1=green, 2=red, 3=blue",
-        },
-        "channel_2": {
-            "type": int,
-            "default": 0,
-            "min": 0,
-            "max": 3,
-            "description": "Second channel: 0=none, 1=green, 2=red, 3=blue",
-        },
+        # "channel_1": {
+        #     "type": int,
+        #     "default": 0,
+        #     "min": 0,
+        #     "max": 3,
+        #     "description": "First channel: 0=grayscale, 1=green, 2=red, 3=blue",
+        # },
+        # "channel_2": {
+        #     "type": int,
+        #     "default": 0,
+        #     "min": 0,
+        #     "max": 3,
+        #     "description": "Second channel: 0=none, 1=green, 2=red, 3=blue",
+        # },
         "flow_threshold": {
             "type": float,
             "default": 0.4,
@@ -214,6 +215,27 @@ def run_cellpose(
             "max": 0.9,
             "description": "Flow threshold for Cellpose segmentation",
         },
+        "cellprob_threshold": {
+            "type": float,
+            "default": 0.0,
+            "min": -6.0,
+            "max": 6.0,
+            "description": "Cell probability threshold (Cellpose 4 parameter)",
+        },
+        "tile_norm_blocksize": {
+            "type": int,
+            "default": 128,
+            "min": 32,
+            "max": 512,
+            "description": "Block size for tile normalization (Cellpose 4 parameter)",
+        },
+        "batch_size": {
+            "type": int,
+            "default": 32,
+            "min": 1,
+            "max": 128,
+            "description": "Batch size for processing multiple images/slices at once",
+        },
         "force_dedicated_env": {
             "type": bool,
             "default": False,
@@ -223,20 +245,21 @@ def run_cellpose(
 )
 def cellpose_segmentation(
     image: np.ndarray,
-    model_type: str = "cyto3",
-    diameter: float = 40.0,
     dim_order: str = "YX",
     channel_1: int = 0,
     channel_2: int = 0,
     flow_threshold: float = 0.4,
+    cellprob_threshold: float = 0.0,
+    tile_norm_blocksize: int = 128,
+    batch_size: int = 32,
     force_dedicated_env: bool = False,
 ) -> np.ndarray:
     """
-    Run Cellpose segmentation on an image.
+    Run Cellpose 4 (Cellpose-SAM) segmentation on an image.
 
     This function takes an image and performs automatic instance segmentation using
-    Cellpose. It supports both 2D and 3D images, various dimension orders, and handles
-    time series data.
+    Cellpose 4 with improved generalization for cellular segmentation. It supports
+    both 2D and 3D images, various dimension orders, and handles time series data.
 
     If Cellpose is not available in the current environment, a dedicated virtual
     environment will be created to run Cellpose.
@@ -245,10 +268,6 @@ def cellpose_segmentation(
     -----------
     image : numpy.ndarray
         Input image
-    model_type : str
-        Cellpose model type: 'cyto'/'cyto2'/'cyto3' for cells, 'nuclei' for nuclei (default: "cyto3")
-    diameter : float
-        Approximate diameter of objects to segment in pixels (default: 40.0)
     dim_order : str
         Dimension order of the input (e.g., 'YX', 'ZYX', 'TZYX') (default: "YX")
     channel_1 : int
@@ -257,6 +276,12 @@ def cellpose_segmentation(
         Second channel: 0=none, 1=green, 2=red, 3=blue (default: 0)
     flow_threshold : float
         Flow threshold for Cellpose segmentation (default: 0.4)
+    cellprob_threshold : float
+        Cell probability threshold (Cellpose 4 parameter) (default: 0.0)
+    tile_norm_blocksize : int
+        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)
 
@@ -266,16 +291,8 @@ def cellpose_segmentation(
         Segmented image with instance labels
     """
     # Convert channel parameters to Cellpose channels list
-    channels = [channel_1, channel_2]
-
-    # Validate parameters
-    valid_models = ["cyto", "cyto2", "cyto3", "nuclei"]
-    if model_type not in valid_models:
-        raise ValueError(
-            f"Invalid model_type: {model_type}. "
-            f"Must be one of: {', '.join(valid_models)}"
-        )
-
+    # 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
 
@@ -293,29 +310,30 @@ def cellpose_segmentation(
         # Prepare arguments for the Cellpose function
         args = {
             "image": image,
-            "model_type": model_type,
-            "diameter": diameter,
             "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,
         }
 
         # Run Cellpose in the dedicated environment
-        print(f"Running Cellpose ({model_type}) in 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(f"Running Cellpose ({model_type}) in current environment...")
+        print("Running Cellpose model in current environment...")
         # Initialize Cellpose model in current environment
-        model = models.Cellpose(gpu=USE_GPU, model_type=model_type)
+        model = models.CellposeModel(gpu=USE_GPU)
 
     # Print status information
     gpu_status = "GPU" if USE_GPU else "CPU"
-    print(f"Using Cellpose {model_type} model on {gpu_status}")
+    print(f"Using Cellpose on {gpu_status}")
     print(
         f"Processing image with shape {image.shape}, dimension order: {dim_order}"
     )
@@ -323,7 +341,14 @@ def cellpose_segmentation(
     # Run segmentation
     try:
         result = run_cellpose(
-            image, model, channels, diameter, flow_threshold, dim_order
+            image,
+            model,
+            channels,
+            flow_threshold,
+            cellprob_threshold,
+            dim_order,
+            tile_norm_blocksize=tile_norm_blocksize,
+            batch_size=batch_size,
         )
 
         print(f"Segmentation complete. Found {np.max(result)} objects.")
@@ -338,10 +363,11 @@ def cellpose_segmentation(
             try:
                 args = {
                     "image": image,
-                    "model_type": model_type,
-                    "diameter": diameter,
                     "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,
@@ -362,150 +388,11 @@ def cellpose_segmentation(
             raise
 
 
-# Add a command-line function to run cellpose segmentation
-def run_cellpose_segmentation():
-    """Run Cellpose segmentation from the command line."""
-    import argparse
-
-    from skimage.io import imread
-    from tifffile import imwrite
-    from tqdm import tqdm
-
-    # Parse arguments
-    parser = argparse.ArgumentParser(
-        description="Runs automatic mask generation on images using Cellpose."
-    )
-    parser.add_argument(
-        "--input", type=str, required=True, help="Path to input images."
-    )
-    parser.add_argument(
-        "--diameter", type=float, default=40.0, help="Diameter of objects."
-    )
-    parser.add_argument(
-        "--channels",
-        type=int,
-        nargs="+",
-        default=[0, 0],
-        help="Channels to use.",
-    )
-    parser.add_argument(
-        "--dim_order",
-        type=str,
-        default="ZYX",
-        help="Dimension order of the input images.",
-    )
-    parser.add_argument(
-        "--model_type",
-        type=str,
-        default="cyto3",
-        choices=["cyto", "cyto2", "cyto3", "nuclei"],
-        help="Model type: 'cyto'/'cyto2'/'cyto3' for cells, 'nuclei' for nuclei",
-    )
-    parser.add_argument(
-        "--flow_threshold",
-        type=float,
-        default=0.4,
-        help="Flow threshold for Cellpose (default: 0.4)",
-    )
-    parser.add_argument(
-        "--max_pixels",
-        type=int,
-        default=4000000,
-        help="Maximum number of pixels to process for 2D images (default: 4000000)",
-    )
-
-    args = parser.parse_args()
-
-    # Validate input folder
-    input_folder = args.input
-    if not os.path.isdir(input_folder):
-        print(
-            f"Error: The input folder '{input_folder}' does not exist or is not accessible."
-        )
-        return 1
-
-    # Find input files
-    input_files = [
-        f
-        for f in os.listdir(input_folder)
-        if f.endswith(".tif") and not f.endswith("_labels.tif")
-    ]
-
-    if not input_files:
-        print(f"No .tif files found in {input_folder}")
-        return 1
-
-    print(f"Found {len(input_files)} files to process")
-
-    # Check if Cellpose is available
-    if not CELLPOSE_AVAILABLE:
-        print(
-            "Error: Cellpose is not installed. Please install it with: pip install cellpose"
-        )
-        return 1
-
-    # Initialize model
-    model = models.Cellpose(gpu=USE_GPU, model_type=args.model_type)
-
-    # Print status
-    gpu_status = "GPU" if USE_GPU else "CPU"
-    print(f"Using Cellpose {args.model_type} model on {gpu_status}")
-
-    # Process each file
-    for input_file in tqdm(input_files, desc="Processing images"):
-        try:
-            # Check image size
-            img = imread(os.path.join(input_folder, input_file))
-
-            if len(img.shape) == 2 or (
-                len(img.shape) == 3 and "C" in args.dim_order
-            ):
-                # For 2D images (potentially with channels)
-                height, width = img.shape[:2]
-                total_pixels = height * width
-                if total_pixels > args.max_pixels:
-                    print(
-                        f"Skipping {input_file} as it exceeds the maximum size of {args.max_pixels} pixels."
-                    )
-                    continue
-
-            print(
-                "\nCheck if image shape corresponds to the dim order that you have given:"
-            )
-            print(
-                f"Image shape: {img.shape}, dimension order: {args.dim_order}\n"
-            )
-
-            # Run segmentation
-            result = run_cellpose(
-                img,
-                model,
-                args.channels,
-                args.diameter,
-                args.flow_threshold,
-                args.dim_order,
-                args.max_pixels,
-            )
-
-            # Save result
-            output_file = os.path.join(
-                input_folder, input_file.replace(".tif", "_labels.tif")
-            )
-            imwrite(output_file, result, compression="zlib")
-
-            print(
-                f"Saved segmentation with {np.max(result)} objects to {output_file}"
-            )
-
-        except (Exception, MemoryError) as e:
-            print(f"Error processing {input_file}: {str(e)}")
-
-    print("\nProcessing complete.")
-    return 0
-
-
-# Run the command-line function if this script is run directly
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(run_cellpose_segmentation())
+# Update cellpose_env_manager.py to install Cellpose 4
+def update_cellpose_env_manager():
+    """
+    Update the cellpose_env_manager to install Cellpose 4
+    """
+    # This function can be called to update the environment manager code
+    # For example, by modifying the pip install command to install the latest version
+    # or specify Cellpose 4 explicitly