From 852f0ae7db01a4776a5ea16832666949b0e8b6f8 Mon Sep 17 00:00:00 2001
From: Xin Bo Qi <z003v5ns@code.siemens.com>
Date: Tue, 9 Mar 2021 15:48:15 +0800
Subject: [PATCH 1/3] add probablity threshold

---
 cellpose2msk/flow2msk.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/cellpose2msk/flow2msk.py b/cellpose2msk/flow2msk.py
index df62ef4..cde6912 100644
--- a/cellpose2msk/flow2msk.py
+++ b/cellpose2msk/flow2msk.py
@@ -5,10 +5,20 @@
 import numpy as np
 from scipy import ndimage as ndimg
 
-def flow2msk(flow, prob, grad=1.0, area=150, volume=500):
+def flow2msk(flow, cell_prob, prob=0.0, grad=0.4, area=20, volume=100):
     shp, dim = flow.shape[:-1], flow.ndim - 1
     l = np.linalg.norm(flow, axis=-1)
-    flow /= l.reshape(shp+(1,));flow[l<grad] = 0
+    flow /= l.reshape(shp+(1,))
+    flow_copy = flow.copy()
+
+    # cell probablity threshold	
+    sigmoid = 1/(1+np.exp(-1*cell_prob))
+    flow[sigmoid<prob] = 0
+
+    # flow magnitude threshold
+    mag = abs(np.amax(flow_copy, axis=-1)/5.0)
+    flow[mag > grad] = flow_copy[mag > grad]
+
     ss = ((slice(None),) * (dim) + ([0,-1],)) * 2
     for i in range(dim):flow[ss[dim-i:-i-2]+(i,)]=0
     sn = np.sign(flow); sn *= 0.5; flow += sn;
@@ -46,7 +56,7 @@ def flow2msk(flow, prob, grad=1.0, area=150, volume=500):
     mask, flow, style, diam = model.eval(
         img, diameter=30, rescale=None, channels=[0,0])
     start = time()
-    water, core, msk = flow2msk(flow[1].transpose(1,2,0), None, 1.0, 20, 100)
+    water, core, msk = flow2msk(flow[1].transpose(1,2,0), flow[2])
     print('flow to mask cost:', time()-start)
     ax1, ax2, ax3, ax4, ax5, ax6 =\
         [plt.subplot(230+i) for i in (1,2,3,4,5,6)]

From 0c2cd1a7d5823d86dea68709ec7c1c5095448718 Mon Sep 17 00:00:00 2001
From: Xin Bo Qi <z003v5ns@code.siemens.com>
Date: Sat, 13 Mar 2021 13:48:25 +0800
Subject: [PATCH 2/3] default prob changed

---
 cellpose2msk/flow2msk.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cellpose2msk/flow2msk.py b/cellpose2msk/flow2msk.py
index cde6912..5aec1ad 100644
--- a/cellpose2msk/flow2msk.py
+++ b/cellpose2msk/flow2msk.py
@@ -5,7 +5,7 @@
 import numpy as np
 from scipy import ndimage as ndimg
 
-def flow2msk(flow, cell_prob, prob=0.0, grad=0.4, area=20, volume=100):
+def flow2msk(flow, cell_prob, prob=0.5, grad=0.4, area=20, volume=100):
     shp, dim = flow.shape[:-1], flow.ndim - 1
     l = np.linalg.norm(flow, axis=-1)
     flow /= l.reshape(shp+(1,))

From fd36c8160cc32ecde8e60a6115533a58c942f357 Mon Sep 17 00:00:00 2001
From: Xin Bo Qi <z003v5ns@code.siemens.com>
Date: Sun, 14 Mar 2021 15:26:37 +0800
Subject: [PATCH 3/3] update for the two thresholds

---
 cellpose2msk/flow2msk.py | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/cellpose2msk/flow2msk.py b/cellpose2msk/flow2msk.py
index 5aec1ad..91d3eb3 100644
--- a/cellpose2msk/flow2msk.py
+++ b/cellpose2msk/flow2msk.py
@@ -5,20 +5,23 @@
 import numpy as np
 from scipy import ndimage as ndimg
 
-def flow2msk(flow, cell_prob, prob=0.5, grad=0.4, area=20, volume=100):
+def estimate_volumes(arr, sigma=3):
+    msk = arr > 50; 
+    idx = np.arange(len(arr), dtype=np.uint32)
+    idx, arr = idx[msk], arr[msk]
+    for k in np.linspace(5, sigma, 5):
+       std = arr.std()
+       dif = np.abs(arr - arr.mean())
+       msk = dif < std * k
+       idx, arr = idx[msk], arr[msk]
+    return arr.mean(), arr.std()
+
+def flow2msk(flow_origin, prob, level=0.5, grad=0.5, area=None, volume=None):
+    flow = flow_origin.copy()
     shp, dim = flow.shape[:-1], flow.ndim - 1
     l = np.linalg.norm(flow, axis=-1)
-    flow /= l.reshape(shp+(1,))
-    flow_copy = flow.copy()
-
-    # cell probablity threshold	
-    sigmoid = 1/(1+np.exp(-1*cell_prob))
-    flow[sigmoid<prob] = 0
-
-    # flow magnitude threshold
-    mag = abs(np.amax(flow_copy, axis=-1)/5.0)
-    flow[mag > grad] = flow_copy[mag > grad]
-
+    flow = flow/l.reshape(shp+(1,))
+    flow[(prob<-np.log(1/(level+1.0e-10)-1))|(l<grad)] = 0
     ss = ((slice(None),) * (dim) + ([0,-1],)) * 2
     for i in range(dim):flow[ss[dim-i:-i-2]+(i,)]=0
     sn = np.sign(flow); sn *= 0.5; flow += sn;
@@ -32,6 +35,9 @@ def flow2msk(flow, cell_prob, prob=0.5, grad=0.4, area=20, volume=100):
     lab, n = ndimg.label(hist, np.ones((3,)*dim))
     volumes = ndimg.sum(hist, lab, np.arange(n+1))
     areas = np.bincount(lab.ravel())
+    mean, std = estimate_volumes(volumes, 2)
+    if volume is None: volume = max(mean-std*3, 50)
+    if area is None: area = volumes // 3
     msk = (areas<area) & (volumes>volume)
     lut = np.zeros(n+1, np.uint32)
     lut[msk] = np.arange(1, msk.sum()+1)