Upgrade cellpose component

src/methods/cellpose/config.vsh.yaml

@@ -1,17 +1,28 @@
+__merge__: /src/api/comp_method.yaml 
+
 name: cellpose
 label: "Cellpose"
-# TODO: update the summary, description and links
-summary: "Cellpose-SAM: cell and nucleus segmentation with superhuman generalization."
-description: "cellpose is an anatomical segmentation algorithm written in Python 3."
-links: # these should point to the documentation of the method
+
+# summary and description are generated by claude sonnet 4.6 based on the paper, github repo and documentation
+summary: Uses predicted flow fields to group pixels into cells via a center-convergence mechanism.
+description: |
+  Rather than detecting cell boundaries directly, Cellpose trains a deep neural network to predict, for every
+  pixel in an image, a vector pointing toward the center of the cell that pixel belongs to, along with a
+  probability map indicating whether each pixel is part of any cell at all. These per-pixel vectors
+  collectively form a "flow field." At inference time, each pixel is treated as a particle that iteratively
+  follows the predicted flow directions, moving step by step toward its eventual destination. Pixels belonging
+  to the same cell all converge to the same fixed point, and any group of pixels that lands at the same fixed
+  point gets assigned a shared cell mask. This center-convergence framing lets the method handle cells of
+  highly varied shapes and sizes without needing to explicitly model cell outlines. The underlying network
+  uses a U-Net-style architecture with skip connections and was trained on a diverse dataset of over 70,000
+  manually segmented objects spanning many cell types and imaging modalities, giving it broad out-of-the-box
+  generalization without requiring retraining or manual parameter tuning for new datasets.
+links:
   documentation: "https://cellpose.readthedocs.io/en/latest/"
   repository: "https://github.com/mouseland/cellpose"
 references:
   doi: "10.1038/s41592-020-01018-x"
 
-
-__merge__: /src/api/comp_method.yaml 
-
 arguments:
   - name: --diameter
     type: double
@@ -49,15 +60,13 @@ resources:
 
 engines:
   - type: docker
-    #image: openproblems/base_pytorch_nvidia:1 # TODO: ideally get gpu image to work
-    image: openproblems/base_python:1
+    image: openproblems/base_pytorch_nvidia:1
     setup:
       - type: python
         pypi: cellpose
       - type: python
         script: from cellpose.models import CellposeModel; model = CellposeModel()
-    __merge__: 
-      - /src/base/setup_spatialdata_partial.yaml
+    __merge__: /src/base/setup_spatialdata_partial.yaml
   - type: native
 
 runners:

src/methods/cellpose/script.py

@@ -1,12 +1,16 @@
 import anndata as ad
 import numpy as np
 import os
-import pandas as pd
 import shutil
 import spatialdata as sd
 import xarray as xr
 from cellpose.models import CellposeModel
 from spatialdata.models import Labels2DModel
+import torch
+
+# Check whether a GPU is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print('Using device:', device, flush=True)
 
 ## VIASH START
 par = {
@@ -38,10 +42,10 @@ def convert_to_lower_dtype(arr):
 transformation = sdata['morphology_mip']['scale0'].image.transform.copy()
 
 print('Initializing Cellpose model', flush=True)
-model = CellposeModel()
+model = CellposeModel(gpu=torch.cuda.is_available())
 
 eval_params = {k: par[k] for k in ("diameter", "flow_threshold", "niter", "min_size", "resample") if par.get(k) is not None}
-print(f"Running Cellpose segmentation with parameters: {eval_params}")
+print('Running Cellpose segmentation with parameters:', eval_params, flush=True)
 masks, _, _ = model.eval(image[0], progress=True, **eval_params)
 
 print('Cellpose segmentation finished, post-processing results', flush=True)