modify track_neck_cut. working for a few cells now

celltk/track_operation.py

@@ -103,7 +103,6 @@ def run_lap(img0, img1, labels0, labels1, DISPLACEMENT=30, MASSTHRES=0.2):
     # This part will make sure that cells outside of these range do not get connected.
     binary_cost[(np.abs(massdiff) > MASSTHRES)] = False
     binary_cost[(dist > DISPLACEMENT)] = False
-
     gp, gc = np.where(binary_cost)
     idx0, idx1 = list(gp), list(gc)
 
@@ -121,7 +120,7 @@ def run_lap(img0, img1, labels0, labels1, DISPLACEMENT=30, MASSTHRES=0.2):
 
 
 def track_neck_cut(img0, img1, labels0, labels1, DISPLACEMENT=10, MASSTHRES=0.2,
-                   EDGELEN=5, THRES_ANGLE=180, WSLIMIT=False):
+                   EDGELEN=5, THRES_ANGLE=180, WSLIMIT=False, SMALL_RAD=3, CANDS_LIMIT=300):
     """
     Adaptive segmentation by using tracking informaiton.
     Separate two objects by making a cut at the deflection. For each points on the outline,
@@ -133,18 +132,20 @@ def track_neck_cut(img0, img1, labels0, labels1, DISPLACEMENT=10, MASSTHRES=0.2,
     THRES_ANGLE (int):  Define the neck points if a triangle has more than this angle.
     STEPLIM (int):      points of neck needs to be separated by at least STEPLIM in parimeters.
     WSLIMIT (bool):     Limit search points to ones overlapped with watershed transformed images. Set it True if calculation is slow.
+
+    SMALL_RAD (int or None): The smallest radius of candidate objects. If you have many cells, set it to None will infer the radius from previous frame.
+    CANDS_LIMIT(int): use lower if slow. limit a number of searches.
+
     """
     labels0, labels = nn_closer(img0, img1, labels0, labels1, DISPLACEMENT, MASSTHRES)
     labels1 = -labels.copy()
 
-    CANDS_LIMIT = 300
-
-    if not hasattr(holder, "SMALL_RAD") and not hasattr(holder, "LARGE_RAD"):
+    if SMALL_RAD is None and not hasattr(holder, 'SMALL_RAD'):
         tracked_area = [i.area for i in regionprops(labels)]
-        holder.LARGE_RAD = np.sqrt(max(tracked_area)/np.pi)
-        holder.SMALL_RAD = np.sqrt(np.percentile(tracked_area, 5)/np.pi)
+        holder.SMALL_RAD = np.sqrt(np.percentile(tracked_area, 5)/np.pi)        
+    elif SMALL_RAD is not None:
+        holder.SMALL_RAD = SMALL_RAD
     SMALL_RAD = holder.SMALL_RAD
-    LARGE_RAD = holder.LARGE_RAD
 
     rps0 = regionprops(labels0, img0)
     unique_labels = np.unique(labels1)
@@ -160,18 +161,19 @@ def track_neck_cut(img0, img1, labels0, labels1, DISPLACEMENT=10, MASSTHRES=0.2,
         if label_id == 0:
             continue
         cc = CellCutter(labels1 == label_id, img1, wlines, small_rad=SMALL_RAD,
-                        large_rad=LARGE_RAD, EDGELEN=EDGELEN, THRES=THRES_ANGLE)
+                        EDGELEN=EDGELEN, THRES=THRES_ANGLE, CANDS_LIMIT=CANDS_LIMIT)
         cc.prepare_coords_set()
         candidates = cc.search_cut_candidates(cc.bw.copy(), cc.coords_set[:CANDS_LIMIT])
         for c in candidates:
             c.raw_label = label_id
         store.append(candidates)
         coords_store.append(cc.coords_set)
+
     coords_store = [i for i in coords_store if i]
     # Attempt a first cut.
     good_cells = _find_best_neck_cut(rps0, store, DISPLACEMENT, MASSTHRES)
     labels0, labels = _update_labels_neck_cut(labels0, labels1, good_cells)
-
+    labels0, labels = nn_closer(img0, img1, labels0, -labels, DISPLACEMENT, MASSTHRES)
     # iteration from here.
     while good_cells:
         rps0 = regionprops(labels0, img0)
@@ -195,6 +197,6 @@ def track_neck_cut(img0, img1, labels0, labels1, DISPLACEMENT=10, MASSTHRES=0.2,
             coords_store.append(coords_set)
         good_cells = _find_best_neck_cut(rps0, store, DISPLACEMENT, MASSTHRES)
         labels0, labels = _update_labels_neck_cut(labels0, labels1, good_cells)
-    labels0, labels = nn_closer(img0, img1, labels0, -labels, DISPLACEMENT, MASSTHRES)
+        labels0, labels = nn_closer(img0, img1, labels0, -labels, DISPLACEMENT, MASSTHRES)
     return labels0, labels
 

celltk/utils/concave_seg.py

@@ -216,23 +216,28 @@ def cellfilter(cell, small_area, large_area, major_minor=2.0):
         return False
 
 
+
 class CellCutter(object):
-    def __init__(self, bw, img, wlines, small_rad, large_rad, EDGELEN=10, THRES=180):
-        DISTLIM_SCALE = 1.5
+    def __init__(self, bw, img, wlines, small_rad, large_rad=0, EDGELEN=10, THRES=180, CANDS_LIMIT=300):
+        """
+        Use lower values for CANDS_LIMIT to speed up
+        """
+        large_rad = 0  # fix it later.
+
         self.bw = bw
         self.img = img
         self.wlines = wlines
         self.EDGELEN = EDGELEN
         self.THRES = THRES
         self.filfunc = partial(cellfilter, small_area=np.pi*small_rad**2, large_area=np.pi*large_rad**2, major_minor=2.0)
         self.goodcells = []
-        self.large_rad = large_rad
+        # self.large_rad = large_rad
         self.small_rad = small_rad
-        self.STEPLIM = self.large_rad * np.pi/2
-        self.DISTLIM = self.large_rad * DISTLIM_SCALE
+        self.STEPLIM = self.small_rad * np.pi/2
         self.coords_set = []
         self.cut_coords = []
         self.cut_cells = []
+        self.CANDS_LIMIT = CANDS_LIMIT
 
     def extract_cell_outlines(self):
         o_coords = find_oriented_coords(self.bw)
@@ -300,18 +305,30 @@ def prepare_coords_set(self):
         coords, _ = CoordsConcaveWs(self.cell.o_coords, wlines=self.wlines, edgelen=self.EDGELEN, thres=self.THRES).run()
         if not coords.any():
             return
-        self.coords_set = self.filter_coords_by_dist(coords)
+        self.coords_set = self.make_sort_coordsset_by_dist(coords)
         if not self.coords_set:
             return
+        self.coords_set = self.coords_set[:self.CANDS_LIMIT]
         self.coords_set = self.filter_coords_by_step(self.cell.o_coords, self.coords_set)        
 
-    def filter_coords_by_dist(self, coords):
+    def make_sort_coordsset_by_dist(self, coords):
         dist = cdist(coords, coords)
-        boolarr = np.triu(dist < self.DISTLIM, 1)
+
+        boolarr = np.triu(np.ones(dist.shape, bool), 1)
         a1 = np.tile(range(len(coords)), len(coords)).reshape((len(coords), len(coords)))
         x1, x2 = a1[boolarr], a1.T[boolarr]
+
+        sortorder = np.argsort(dist[boolarr])
+        x1, x2 = x1[sortorder], x2[sortorder]
         return [[coords[i1], coords[i2]] for i1, i2 in zip(x1, x2)]
 
+    # def filter_coords_by_dist(self, coords):
+    #     dist = cdist(coords, coords)
+        # boolarr = np.triu(dist < self.DISTLIM, 1)
+    #     a1 = np.tile(range(len(coords)), len(coords)).reshape((len(coords), len(coords)))
+    #     x1, x2 = a1[boolarr], a1.T[boolarr]
+    #     return [[coords[i1], coords[i2]] for i1, i2 in zip(x1, x2)]
+
     def filter_coords_by_step(self, coords, coords_set):
         new_coords_set = []
         for c0, c1 in coords_set: