From ff42665854d9556b421c78e0c49236586f8cca51 Mon Sep 17 00:00:00 2001
From: Marcel <46164444+MSt-10@users.noreply.github.com>
Date: Mon, 31 Mar 2025 16:16:28 +0200
Subject: [PATCH] Fixed error if no cells detected, return type more consistent

- Return type of model.predict(...) was changed here: https://github.com/vanvalenlab/cellSAM/commit/73bd4e3c5c96935bb854c4910bc49471c56bb3b9#diff-f21d549921220b19fc1455312955b96866ecbf013502dc5e85974d1995559e1fR303, caused error if no cells were detected.
- In case of no cells, returned values consistent (single channel uint8 mask, empty bbox list)
---
 cellSAM/model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cellSAM/model.py b/cellSAM/model.py
index 23283be..95bbe9c 100644
--- a/cellSAM/model.py
+++ b/cellSAM/model.py
@@ -144,9 +144,9 @@ def segment_cellular_image(
         model, img = model.to(device), img.to(device)
 
     preds = model.predict(img, x=None, boxes_per_heatmap=bounding_boxes, device=device, fast=fast)
-    if preds is None:
+    if preds[0] is None:
         warn("No cells detected in the image.")
-        return np.zeros(img.shape[1:], dtype=np.int32), None, None
+        return np.zeros(img.shape[-2:], dtype=np.uint8), None, torch.empty((1, 4))
 
     segmentation_predictions, _, x, bounding_boxes = preds