Merge pull request #1184 from rishant3441/cellpose4-support

Adds Cellpose-SAM Support. Verified that it fixes both #1180 and #1178.

Author: chriski777

README.md

@@ -50,6 +50,8 @@ Pachitariu, M., Stringer, C., Schröder, S., Dipoppa, M., Rossi, L. F., Carandin
 
 ## Local installation
 
+If you are using a GPU, make sure its drivers and the cuda libraries are correctly installed.
+
 1. Install an [Anaconda](https://www.anaconda.com/download/) distribution of Python -- Choose **Python 3.8** and your operating system. Note you might need to use an anaconda prompt if you did not add anaconda to the path.
 2. Open an anaconda prompt / command prompt with `conda` for **python 3** in the path
 3. Create a new environment with `conda create --name suite2p python=3.9`.
@@ -76,8 +78,28 @@ This package relies on the awesomeness of [pyqtgraph](http://pyqtgraph.org/), [P
 
 The software has been heavily tested on Windows 10 and Ubuntu 18.04, and less well tested on Mac OS. Please post an [issue](https://github.com/MouseLand/suite2p/issues) if you have installation problems. 
 
+### GPU version (CUDA) on Windows or Linux
+
+If you plan on running Cellpose-SAM (anatomical ROI detection), you may want to install a GPU version of *torch*. To use your NVIDIA GPU with python, you will need to make sure the NVIDIA driver for your GPU is installed, check out this [website](https://www.nvidia.com/Download/index.aspx?lang=en-us) to download it. You can also install the CUDA toolkit, or use the pytorch cudatoolkit (installed below with conda). If you have trouble with the below install, we recommend installing the CUDA toolkit yourself, choosing one of the 11.x releases [here](https://developer.nvidia.com/cuda-toolkit-archive).
+
+With the latest versions of pytorch on Linux, as long as the NVIDIA drivers are installed, the GPU version is installed by default with pip. You can check if the GPU support is working by opening the GUI. If the GPU is working then the `GPU` box will be checked and the `CUDA` version will be displayed in the command line. 
+
+If it's not working, we will need to remove the CPU version of torch:
+~~~
+pip uninstall torch
+~~~
+
+To install the GPU version of torch, follow the instructions [here](https://pytorch.org/get-started/locally/). The pip or conda installs should work across platforms, you will need torch and torchvision, e.g. for windows + cuda 12.6 the command is
+~~~
+pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu126
+~~~
+
+Info on how to install several older versions is available [here](https://pytorch.org/get-started/previous-versions/). After install you can check `conda list` for `pytorch`, and its version info should have `cuXX.X`, not `cpu`.
+
 ### Installing the latest github version of the code
 
+If you are using a GPU, make sure its drivers and the cuda libraries are correctly installed.
+
 The simplest way is
 ~~~
 pip install git+https://github.com/MouseLand/suite2p.git
@@ -90,6 +112,8 @@ If you want to download and edit the code, and use that version,
 
 ### Installation for developers
 
+If you are using a GPU, make sure its drivers and the cuda libraries are correctly installed.
+
 1. Clone the repository and `cd suite2p` in an anaconda prompt / command prompt with `conda` for **python 3** in the path
 2. Run `conda create --name suite2p python=3.9`
 3. To activate this new environment, run `conda activate suite2p` (you will have to activate every time you want to run suite2p)

setup.py

@@ -10,7 +10,7 @@
         "matplotlib",
         "scipy>=1.9.0",
         "scikit-learn",
-        "cellpose<=3.1.1.2",
+        "cellpose>=4.0.1",
         "scanimage-tiff-reader>=1.4.1"
         ]
 

suite2p/default_ops.py

@@ -130,10 +130,10 @@ def default_ops():
             0,  # run cellpose to get masks on 1: max_proj / mean_img; 2: mean_img; 3: mean_img enhanced, 4: max_proj
         "diameter": 0,  # use diameter for cellpose, if 0 estimate diameter
         "cellprob_threshold": 0.0,  # cellprob_threshold for cellpose
-        "flow_threshold": 1.5,  # flow_threshold for cellpose
+        "flow_threshold": 0.4,  # flow_threshold for cellpose
         "spatial_hp_cp": 0,  # high-pass image spatially by a multiple of the diameter
         "pretrained_model":
-            "cyto",  # path to pretrained model or model type string in Cellpose (can be user model)
+            "cpsam",  # path to pretrained model or model type string in Cellpose (can be user model)
 
         # classification parameters
         "soma_crop":

suite2p/detection/anatomical.py

@@ -4,8 +4,8 @@
 import numpy as np
 from typing import Any, Dict
 from scipy.ndimage import find_objects, gaussian_filter
-from cellpose.models import CellposeModel, Cellpose
-from cellpose import transforms, dynamics
+from cellpose.models import CellposeModel
+from cellpose import transforms, dynamics, core
 from cellpose.utils import fill_holes_and_remove_small_masks
 from cellpose.transforms import normalize99
 import time
@@ -34,7 +34,7 @@ def patch_detect(patches, diam):
     print("refining masks using cellpose")
     npatches = len(patches)
     ly = patches[0].shape[0]
-    model = Cellpose()
+    model = CellposeModel(gpu=True if core.use_gpu() else False)
     imgs = np.zeros((npatches, ly, ly, 2), np.float32)
     for i, m in enumerate(patches):
         imgs[i, :, :, 0] = transforms.normalize99(m)
@@ -46,7 +46,7 @@ def patch_detect(patches, diam):
     batch_size = 8 * 224 // ly
     tic = time.time()
     for j in np.arange(0, npatches, batch_size):
-        y = model.cp.network(imgs[j:j + batch_size])[0]
+        y = model.net(imgs[j:j + batch_size])[0]
         y = y[:, :, ysub[0]:ysub[-1] + 1, xsub[0]:xsub[-1] + 1]
         y = y.asnumpy()
         for i, yi in enumerate(y):
@@ -99,14 +99,13 @@ def refine_masks(stats, patches, seeds, diam, Lyc, Lxc):
     return stats
 
 
-def roi_detect(mproj, diameter=None, cellprob_threshold=0.0, flow_threshold=1.5,
+def roi_detect(mproj, diameter=None, cellprob_threshold=0.0, flow_threshold=0.4,
                pretrained_model=None):
-    pretrained_model = "cyto3" if pretrained_model is None else pretrained_model
-    if not os.path.exists(pretrained_model):
-        model = Cellpose(model_type=pretrained_model)
-    else:
-        model = CellposeModel(pretrained_model=pretrained_model)
-    masks = model.eval(mproj, channels=[0, 0], diameter=diameter,
+    if diameter == 0:
+        diameter = None
+    pretrained_model = "cpsam" if pretrained_model is None else pretrained_model
+    model = CellposeModel(pretrained_model=pretrained_model, gpu=True if core.use_gpu() else False)
+    masks = model.eval(mproj, diameter=diameter,
                        cellprob_threshold=cellprob_threshold,
                        flow_threshold=flow_threshold)[0]
     shape = masks.shape

suite2p/detection/chan2detect.py

@@ -5,6 +5,7 @@
 from scipy.ndimage import gaussian_filter
 from ..extraction import masks
 from . import utils
+import traceback
 """
 identify cells with channel 2 brightness (aka red cells)
 
@@ -90,7 +91,10 @@ def cellpose_overlap(stats, mimg2):
         ypix0, xpix0 = stats[i]["ypix"], stats[i]["xpix"]
         smask[ypix0, xpix0] = 1
         ious = utils.mask_ious(masks, smask)[0]
-        iou = ious.max()
+        if ious.size > 0:
+            iou = ious.max()
+        else:
+            iou = 0.0
         redstats[
             i,
         ] = np.array([iou > 0.25, iou])  #this had the wrong dimension
@@ -112,10 +116,11 @@ def detect(ops, stats):
         try:
             print(">>>> CELLPOSE estimating masks in anatomical channel")
             redstats, masks = cellpose_overlap(stats, mimg2)
-        except:
+        except Exception as e:
             print(
                 "ERROR importing or running cellpose, continuing without anatomical estimates"
             )
+            traceback.print_exc()
 
     if redstats is None:
         redstats = intensity_ratio(ops, stats)