bug with min_train_masks

MouseLand · Sep 14, 2024 · 52f75f9 · 52f75f9
1 parent 10d3554
commit 52f75f9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cellpose/train.py b/cellpose/train.py
@@ -293,6 +293,7 @@ def _process_train_test(train_data=None, train_labels=None, train_files=None,
             if train_probs is not None:
                 train_probs = train_probs[ikeep]
             diam_train = diam_train[ikeep]
+            nimg = len(train_data)
 
     ### normalize probabilities
     train_probs = 1. / nimg * np.ones(nimg,
@@ -410,7 +411,7 @@ def train_seg(net, train_data=None, train_labels=None, train_files=None,
             "channel_axis": channel_axis,
             "rgb": rgb
         }
-
+    
     net.diam_labels.data = torch.Tensor([diam_train.mean()]).to(device)
 
     nimg = len(train_data) if train_data is not None else len(train_files)