small fixes for uint32, colab notebooks, cellprob return

cellpose/dynamics.py

@@ -734,7 +734,10 @@ def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
         mask = utils.fill_holes_and_remove_small_masks(mask, min_size=min_size)
 
     if mask.dtype == np.uint32:
-        dynamics_logger.warning(
+        if mask.max() < 2**16:
+            mask = mask.astype("uint16")
+        else:
+            dynamics_logger.warning(
             "more than 65535 masks in image, masks returned as np.uint32")
 
     return mask

cellpose/io.py

@@ -894,3 +894,5 @@ def save_masks(images, masks, flows, file_names, png=True, tif=False, channels=[
               )
         #save full flow data
         imsave(os.path.join(flowdir, basename + '_dP' + suffix + '.tif'), flows[1])
+        # save cellprob
+        imsave(os.path.join(flowdir, basename + '_cellprob' + suffix + '.tif'), flows[2])

cellpose/train.py

@@ -72,10 +72,10 @@ def _reshape_norm(data, channel_axis=None, normalize_params={"normalize": False}
                 # put channel axis first 
                 td = np.moveaxis(td, channel_axis0, 0)
                 td = td[:3] # keep at most 3 channels
-            if td.ndim == 2 or (td.ndim == 3 and td.shape[0] == 1):
+            if td.ndim == 2:
                 td = np.stack((td, 0*td, 0*td), axis=0)
             elif td.ndim == 3 and td.shape[0] < 3:
-                td = np.concatenate((td, 0*td[:1]), axis=0)
+                td = np.concatenate((td, np.zeros((3 - td.shape[0], *td.shape[1:]), dtype=td.dtype)), axis=0)
             data_new.append(td)
         data = data_new
     if normalize_params["normalize"]:

cellpose/utils.py

@@ -213,12 +213,12 @@ def masks_to_outlines(masks):
         return outlines
 
 
-def outlines_list(masks, multiprocessing_threshold=1000, multiprocessing=None):
+def outlines_list(masks, multiprocessing_threshold=50000, multiprocessing=None):
     """Get outlines of masks as a list to loop over for plotting.
 
     Args:
         masks (ndarray): Array of masks.
-        multiprocessing_threshold (int, optional): Threshold for enabling multiprocessing. Defaults to 1000.
+        multiprocessing_threshold (int, optional): Threshold for enabling multiprocessing. Defaults to 50000.
         multiprocessing (bool, optional): Flag to enable multiprocessing. Defaults to None.
 
     Returns:
@@ -529,6 +529,8 @@ def stitch3D(masks, stitch_threshold=0.25):
     mmax = masks[0].max()
     empty = 0
     for i in trange(len(masks) - 1):
+        if masks.dtype == "uint16" and int(mmax) > max(0, 2**16 - 5 - masks[i + 1].max()):
+            masks = masks.astype("uint32")
         iou = metrics._intersection_over_union(masks[i + 1], masks[i])[1:, 1:]
         if not iou.size and empty == 0:
             masks[i + 1] = masks[i + 1]

notebooks/run_Cellpose-SAM.ipynb

@@ -101,7 +101,8 @@
       },
       "outputs": [],
       "source": [
-        "!pip install git+https://www.github.com/mouseland/cellpose.git"
+        "!uv pip install git+https://www.github.com/mouseland/cellpose.git\n",
+        "!uv pip install git+https://github.com/facebookresearch/dinov3.git"
       ]
     },
     {
@@ -135,7 +136,11 @@
         "if core.use_gpu()==False:\n",
         "  raise ImportError(\"No GPU access, change your runtime\")\n",
         "\n",
-        "model = models.CellposeModel(gpu=True)"
+        "model = models.CellposeModel(gpu=True)\n",
+        "\n",
+        "### You can also use other pretrained models, like the DINO models \n",
+        "# model = models.CellposeModel(gpu=True, pretrained_model=\"cpdino\")\n",
+        "# model = models.CellposeModel(gpu=True, pretrained_model=\"cpdino-vitb\")"
       ]
     },
     {
@@ -225,7 +230,7 @@
         "\n",
         "- If you have a histological image taken in brightfield, you don't need to adjust the channels.\n",
         "\n",
-        "- If you have a fluroescent image with multiple stains, you should choose one channel with a cytoplasm/membrane stain, one channel with a nuclear stain, and set the third channel to `None`. Choosing multiple channels may produce segmentaiton of all the structures in the image. If you have retrained the model on your data with a thrid stain (described below), you can run segmentation with all channels. "
+        "- If you have a fluroescent image with multiple stains, you should choose one channel with a cytoplasm/membrane stain, one channel with a nuclear stain, and set the third channel to `None`. Choosing multiple channels may produce segmentaiton of all the structures in the image. If you have retrained the model on your data with a third stain (described below), you can run segmentation with all channels. "
       ]
     },
     {

notebooks/test_Cellpose-SAM.ipynb

@@ -3,8 +3,8 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "view-in-github"
       },
       "source": [
         "<a href=\"https://colab.research.google.com/github/MouseLand/cellpose/blob/main/notebooks/test_Cellpose-SAM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
@@ -54,7 +54,8 @@
       },
       "outputs": [],
       "source": [
-        "!pip install git+https://www.github.com/mouseland/cellpose.git"
+        "!uv pip install git+https://www.github.com/mouseland/cellpose.git\n",
+        "!uv pip install git+https://github.com/facebookresearch/dinov3.git"
       ]
     },
     {
@@ -87,7 +88,11 @@
         "if core.use_gpu()==False:\n",
         "  raise ImportError(\"No GPU access, change your runtime\")\n",
         "\n",
-        "model = models.CellposeModel(gpu=True)"
+        "model = models.CellposeModel(gpu=True)\n",
+        "\n",
+        "### You can also use other pretrained models, like the DINO models \n",
+        "# model = models.CellposeModel(gpu=True, pretrained_model=\"cpdino\")\n",
+        "# model = models.CellposeModel(gpu=True, pretrained_model=\"cpdino-vitb\")"
       ]
     },
     {
@@ -314,6 +319,11 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fd-6Hji-n9_H"
+      },
+      "outputs": [],
       "source": [
         "# DISPLAY RESULTS stitching\n",
         "plt.figure(figsize=(15,3))\n",
@@ -326,22 +336,17 @@
         "  imgout[outX, outY] = np.array([255,75,75])\n",
         "  plt.imshow(imgout)\n",
         "  plt.title('iplane = %d'%iplane)"
-      ],
-      "metadata": {
-        "id": "fd-6Hji-n9_H"
-      },
-      "execution_count": null,
-      "outputs": []
+      ]
     }
   ],
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "provenance": [],
-      "include_colab_link": true
+      "include_colab_link": true,
+      "provenance": []
     },
     "kernelspec": {
-      "display_name": "cellpose",
+      "display_name": "s2p_test",
       "language": "python",
       "name": "python3"
     },
@@ -355,9 +360,9 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.11.9"
+      "version": "3.11.14"
     }
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}

notebooks/train_Cellpose-SAM.ipynb

@@ -3,8 +3,8 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "view-in-github",
-        "colab_type": "text"
+        "colab_type": "text",
+        "id": "view-in-github"
       },
       "source": [
         "<a href=\"https://colab.research.google.com/github/MouseLand/cellpose/blob/main/notebooks/train_Cellpose-SAM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
@@ -113,7 +113,7 @@
       },
       "outputs": [],
       "source": [
-        "!pip install git+https://www.github.com/mouseland/cellpose.git"
+        "!uv pip install git+https://www.github.com/mouseland/cellpose.git"
       ]
     },
     {
@@ -346,8 +346,8 @@
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "provenance": [],
-      "include_colab_link": true
+      "include_colab_link": true,
+      "provenance": []
     },
     "kernelspec": {
       "display_name": "cellpose",