choice cp3 or cpsam in segment cells

Author: gletort

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
     "scikit-image",
     "lxml",
     "packaging",
-    "cellpose[distributed]",
+    "cellpose[distributed]>3",
 ]
 
 [project.optional-dependencies]

src/fish_feats/NapaCells.py

@@ -198,7 +198,7 @@ def __init__( self, ffeats ):
         super().__init__()
         layout = QVBoxLayout()
 
-        methods = [ "Epyseg", "CellPose", "Load segmented file", "Empty" ]
+        methods = [ "Epyseg", "CellPoseSAM", "CellPose3", "Load segmented file", "Empty" ]
         defmeth = "Epyseg"
         celldiameter = 30
         self.chunksize = 1500
@@ -210,6 +210,9 @@ def __init__( self, ffeats ):
                 celldiameter = int(float(self.paras["cell_diameter"]))
             if "method" in self.paras:
                 defmeth = self.paras["method"]
+                ## compatibility with previous version
+                if defmeth == "CellPose":
+                    defmeth = "CellPoseSAM"
         if self.mig.junction_filename(dim=2, ifexist=True) != "":
             defmeth = "Load segmented file"
         self.junction_filename = self.mig.junction_filename(dim=2, ifexist=True)
@@ -253,7 +256,7 @@ def __init__( self, ffeats ):
 
     def visibility( self ):
         """ Set the visibility of the parameters according to the method """
-        self.diam_group.setVisible( self.methodsChoice.currentText() in ["CellPose"] )
+        self.diam_group.setVisible( self.methodsChoice.currentText() in ["CellPose", "CellPoseSAM", "CellPose3"] )
         self.file_group.setVisible( self.methodsChoice.currentText() == "Load segmented file" )
 
     def get_junction_filename( self ):

src/fish_feats/SegmentObj.py

@@ -214,12 +214,73 @@ def junctions_to_label(img):
 
 def run_cellpose(img, scaleXY, diameter=7, verbose=True):
     from cellpose import models
-    model = models.CellposeModel(gpu=True, model_type='cyto3') 
+    model = models.CellposeModel( gpu=True ) 
     diamet = diameter/scaleXY   ## increase it ?
-    mask, flow, style = model.eval(img, invert=False, diameter=diamet, channels=[0,0], do_3D=False, cellprob_threshold=0.05)
+    mask, flow, style = model.eval(img, invert=False, diameter=diamet, do_3D=False, cellprob_threshold=0.05)
     ## convert cellpose result to label image (cellpose result are not touching border, make it as junctions)
     return fromcellpose_tojunctions(mask)
 
+def run_cellpose_3( img, scaleXY, diameter=7, verbose=True, progress_bar=None ):
+    """ Use CellPose3 instead of main Cellpose(SAM), using Appose environment """
+    feature = "cp3-"+ut.get_cuda_feature()
+    try:
+        pixi_file = resources.files("fish_feats.resources").joinpath("pixi_cp3.toml")
+        ut.show_info("Build/Load tensorflow environment")
+        env = appose.pixi( pixi_file ).log_debug()
+        env = env.subscribe_output( lambda line: print("OUT:", line, end="") )
+        env = env.subscribe_error( lambda line: print("DBG:", line, end="") )
+        env = env.environment( feature ).build()
+        ut.show_info(f"Environment built at: {env.base()}")
+        python = env.python().init("import numpy as np; from cellpose import models" \
+        "")
+        if progress_bar is None:
+            progress_bar = ut.start_progress( None, total=1, descr="Cellpose3 segmentation" )
+            toclose = True
+        else:
+            progress_bar.set_description( "Cellpose3 segmentation" )
+            toclose = False
+        
+        def log_listener(event):
+            """ Transfer appose task message to the main logger """
+            if event.current and event.maximum:
+                print( f"Segmenting slice {event.current}/{event.maximum}" )
+            else:
+                if event.message:
+                    print( f"[task] {event.message} " )
+
+        try:
+            cellpose3_script = resources.files("fish_feats.resources").joinpath("run_cellpose3.py")
+            cellpose3_script = cellpose3_script.read_text()
+            with share_as_ndarray(img) as image:
+                task = python.task( cellpose3_script )
+                task.listen( log_listener )
+                task.inputs["img"] = image 
+                task.inputs["cell_diameter"] = diameter
+                task.inputs["scale_xy"] = scaleXY
+                task.wait_for()
+                result = image.ndarray()
+                mask = np.uint16( result )
+                return fromcellpose_tojunctions(mask)
+        except Exception as e:
+            raise RuntimeError("Running cellpose3 in separated environement failed") from e
+        finally:
+            python.close()
+            if toclose:
+                ut.close_progress( None, progress_bar=progress_bar )
+    except Exception as e:
+        raise RuntimeError("Cellpose3 in separated environement failed") from e
+
+def getNuclei_stardist2DAsso3D(nucimg, scaleXY, proba=0.55, overlap=0.1, assoMode="Munkres", assolim=3, threshold_overlap=0.25, verbose=True, progress_bar=None):
+    """ Segment nuclei with Stardist2D and reconstruct in 3D - return the nuclei list """
+    
+    ## segment 2D
+    appose = not ut.has_dependency( "stardist" )
+    if appose:
+        labnuc = stardist2D_appose(nucimg, prob=proba, over=overlap, progress_bar=progress_bar)
+    else:
+        labnuc = stardist2D_local(nucimg, prob=proba, over=overlap, progress_bar=progress_bar)
+
+
 def fromcellpose_tojunctions(mask):
     bined = binary_closing( find_boundaries(mask), footprint=np.ones((10,10)) )
     return junctions_to_label( skeletonize( bined ) )
@@ -233,9 +294,13 @@ def segmentJunctions(img, mode, scaleXY, imagedir, imagename, diameter=7, chunki
             if os.path.exists(tmpdir_path):
                 os.chmod(tmpdir_path, stat.S_IRUSR|stat.S_IRGRP|stat.S_IROTH|stat.S_IWUSR|stat.S_IWGRP|stat.S_IWOTH|stat.S_IXUSR)
                 shutil.rmtree(tmpdir_path, ignore_errors=True)
-    if (isinstance(mode, str) and mode.lower() == "cellpose"):
-        show_info("Starting segmentation with CellPose...")
-        seged = run_cellpose(img, scaleXY, diameter=diameter, verbose=verbose)
+    if (isinstance(mode, str) and mode[0:8].lower() == "cellpose"):
+        if mode.lower() == "cellposesam":
+            show_info("Starting segmentation with CellPose-SAM...")
+            seged = run_cellpose(img, scaleXY, diameter=diameter, verbose=verbose)
+        elif mode.lower() == "cellpose3":
+            show_info("Starting segmentation with CellPose3, cyto3 model...")
+            seged = run_cellpose_3(img, scaleXY, diameter=diameter, verbose=verbose)
     if (isinstance(mode, str) and mode.lower() == "epyseg-dask"):
         from fish_feats.DaskedEpyseg import run_dasked_epyseg
         show_info("Starting segmentation with dasked Epyseg...")

src/fish_feats/Utils.py

@@ -800,6 +800,21 @@ def get_env_name():
     env_name = "cuda" if is_gpu_platform else "default"
     return env_name
 
+def get_cuda_feature():
+    """ Get the drivers cuda version through nvidia-smi and build the relevant feature for pixi """
+    import subprocess, re
+    result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
+    match = re.search(r'CUDA Version:\s+([\d.]+)', result.stdout)
+    if match:
+        cuda_version = match.group(1)
+
+        ## get only major version
+        if cuda_version.split( "." )[0] == "12":
+            return "cu126"
+        if cuda_version.split( "." )[0] == "13":
+            return "cu130"
+    return "cpu"
+
 #### Handle versions of napari
 def version_napari_above( compare_version ):
     """ Get the current version of napari """

src/fish_feats/resources/pixi.toml

@@ -1,5 +1,5 @@
 [workspace]
-name = "tensorflow-env"
+name = "fishfeats-tensorflow-env"
 version = "0.0.0"
 description = "Modules with tensorflow for FishFeats: epyseg, stardist, tensorflow and keras"
 authors = ["Gaelle Letort"]

src/fish_feats/resources/pixi_cp3.toml

@@ -0,0 +1,90 @@
+[workspace]
+authors = [
+    "Gaelle Letort <gaelle.letort@pasteur.fr>",
+]
+channels = ["conda-forge"]
+name = "fishfeats_cellpose3"
+platforms = ["osx-arm64", "win-64", "linux-64", "osx-64"]
+version = "0.1.0"
+
+[dependencies]
+python = ">=3.10.3,<3.13"
+packaging = ">=26.2,<27"
+
+[pypi-dependencies]
+appose = ">=0.10.1, <0.12"
+
+# --- PyTorch features ---
+# See for more details on torch versions and CUDA support: https://pytorch.org/get-started/previous-versions/
+# macOS does not support CUDA and will always fall back to CPU.
+
+# CPU: all platforms use the CPU index
+[feature.cpu.target.linux.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+[feature.cpu.target.win.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+[feature.cpu.target.osx-arm64.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+# Deprecated: osx-64 support should be drop in favor of osx-arm64
+[feature.cpu.target.osx-64.pypi-dependencies]
+torch = { version = ">=2.2.2, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.17.2, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+# CUDA 12.6: Linux/Windows use cu126 index, macOS falls back to CPU
+[feature.cu126.target.linux.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cu126" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cu126" }
+
+[feature.cu126.target.win.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cu126" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cu126" }
+
+[feature.cu126.target.osx-arm64.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+# Deprecated: osx-64 support should be drop in favor of osx-arm64
+[feature.cu126.target.osx-64.pypi-dependencies]
+torch = { version = ">=2.2.2, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.17.2, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+# CUDA 13.0: Linux/Windows use cu130 index, macOS falls back to CPU
+[feature.cu130.target.linux.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cu130" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cu130" }
+
+[feature.cu130.target.win.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cu130" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cu130" }
+
+[feature.cu130.target.osx-arm64.pypi-dependencies]
+torch = { version = ">=2.9, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.24, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+# Deprecated: osx-64 support should be drop in favor of osx-arm64
+[feature.cu130.target.osx-64.pypi-dependencies]
+torch = { version = ">=2.2.2, <2.11", index = "https://download.pytorch.org/whl/cpu" }
+torchvision = { version = ">=0.17.2, <0.26", index = "https://download.pytorch.org/whl/cpu" }
+
+# --- CellPose version features ---
+[feature.cp3.pypi-dependencies]
+cellpose = ">=3, <4"
+
+[feature.cp4.pypi-dependencies]
+cellpose = ">=4, <5"
+
+# --- Environments ---
+[environments]
+default = { features = ["cp3", "cu126"] }
+cp3-cpu = { features = ["cp3", "cpu"] }
+cp4-cpu = { features = ["cp4", "cpu"] }
+cp3-cu126 = { features = ["cp3", "cu126"] }
+cp4-cu126 = { features = ["cp4", "cu126"] }
+cp3-cu130 = { features = ["cp3", "cu130"] }
+cp4-cu130 = { features = ["cp4", "cu130"] }

src/fish_feats/resources/run_cellpose3.py

@@ -0,0 +1,41 @@
+import logging
+import numpy as np
+appose_mode = 'task' in globals()
+
+if not appose_mode:
+    print("not implemented")
+    exit
+try:
+    from cellpose import models
+except ModuleNotFoundError as e:
+    print("ERR: Problem with Cellpose3 installation")
+    raise e 
+
+class ApposeLogHandler(logging.Handler):
+    def emit(self, record):
+        msg = self.format(record)
+        task.update( message=msg )
+
+def setup_logger( name="" ):
+    logger = logging.getLogger(name)
+    handler = ApposeLogHandler()
+    formatter = logging.Formatter('[FishFeats-Cellpose3]-%(message)s')
+    handler.setFormatter( formatter )
+    logger.addHandler(handler)
+    logger.setLevel( 20 )
+    return logger
+
+logger = setup_logger()    
+logger.info("Starting segmentation")
+
+image = np.array( img.ndarray() )
+logger.info("Starting")
+    
+model = models.CellposeModel( gpu=True, model_type='cyto3' ) 
+diamet = cell_diameter/scale_xy  
+mask, flow, style = model.eval(image, invert=False, diameter=diamet, channels=[0,0], do_3D=False, cellprob_threshold=0.05)
+
+if mask is not None:
+    ## Write the results on the shared memory object
+    img.ndarray()[:] = mask 
+logger.info( "Segmentation finished" )