Update to sgn detection training

Author: constantinpape

flamingo_tools/segmentation/postprocessing.py

@@ -774,7 +774,8 @@ def split_object(object_id):
         nonlocal offset
 
         row = segmentation_table[segmentation_table.label_id == object_id]
-        if min_size and row.n_pixels.values[0] < min_size:
+        if row.n_pixels.values[0] < min_size:
+            # print(object_id, ": min-size")
             return [object_id]
 
         bb_min = np.array([
@@ -788,6 +789,12 @@ def split_object(object_id):
         bb_max = np.minimum(bb_max.astype(int) + 1, np.array(list(segmentation.shape)))
         bb = tuple(slice(mi, ma) for mi, ma in zip(bb_min, bb_max))
 
+        # This is due to segmentation artifacts.
+        bb_shape = bb_max - bb_min
+        if (bb_shape > 500).any():
+            print(object_id, "has a too large shape:", bb_shape)
+            return [object_id]
+
         seg = segmentation[bb]
         mask = ~find_boundaries(seg)
         dist = distance_transform_edt(mask, sampling=resolution)
@@ -798,6 +805,7 @@ def split_object(object_id):
         maxima = peak_local_max(dist, min_distance=3, exclude_border=True)
 
         if len(maxima) == 1:
+            # print(object_id, ": max len")
             return [object_id]
 
         with lock:
@@ -819,14 +827,20 @@ def split_object(object_id):
 
         keep_ids = seg_ids[sizes > min_size]
         if len(keep_ids) < 2:
+            # print(object_id, ": keep-id")
             return [object_id]
 
         elif len(keep_ids) != len(seg_ids):
             new_seg[~np.isin(new_seg, keep_ids)] = 0
             new_seg = watershed(hmap, markers=new_seg, mask=seg_mask)
 
-        output[bb][seg_mask] = new_seg[seg_mask]
-        return seg_ids.tolist()
+        with lock:
+            out = output[bb]
+            out[seg_mask] = new_seg[seg_mask]
+            output[bb] = out
+
+        # print(object_id, ":", len(keep_ids))
+        return keep_ids.tolist()
 
         # import napari
         # v = napari.Viewer()
@@ -839,11 +853,15 @@ def split_object(object_id):
     if component_labels is None:
         object_ids = segmentation_table.label_id.values
     else:
-        object_ids = segmentation_table[segmentation_table.isin(component_labels)].label_id.values
+        object_ids = segmentation_table[segmentation_table.component_labels.isin(component_labels)].label_id.values
 
     if n_threads is None:
         n_threads = mp.cpu_count()
 
+    # new_id_mapping = []
+    # for object_id in tqdm(object_ids, desc="Split non-convex objects"):
+    #     new_id_mapping.append(split_object(object_id))
+
     with futures.ThreadPoolExecutor(n_threads) as tp:
         new_id_mapping = list(
             tqdm(tp.map(split_object, object_ids), total=len(object_ids), desc="Split non-convex objects")

scripts/la-vision/train_sgn_detection.py

@@ -12,7 +12,8 @@
 from utils.training.training import supervised_training  # noqa
 from detection_dataset import DetectionDataset, MinPointSampler  # noqa
 
-ROOT = "./la-vision-sgn-new"  # noqa
+# ROOT = "./la-vision-sgn-new"  # noqa
+ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/SGN/sgn-detection"
 
 TRAIN = os.path.join(ROOT, "images")
 TRAIN_EMPTY = os.path.join(ROOT, "empty_images")
@@ -24,6 +25,7 @@
 def _get_paths(split, train_folder, label_folder, n=None):
     image_paths = sorted(glob(os.path.join(train_folder, "*.tif")))
     label_paths = sorted(glob(os.path.join(label_folder, "*.csv")))
+    assert len(image_paths) > 0
     assert len(image_paths) == len(label_paths)
     if n is not None:
         image_paths, label_paths = image_paths[:n], label_paths[:n]

scripts/synapse_marker_detection/detection_dataset.py

@@ -206,12 +206,12 @@ def _get_sample(self, index):
                     raise RuntimeError(f"Could not sample a valid batch in {self.max_sampling_attempts} attempts")
 
         # For synapse detection.
-        # label = process_labels(coords, shape, self.sigma, self.eps, bb=bb)
+        label = process_labels(coords, shape, self.sigma, self.eps, bb=bb)
 
         # For SGN detection with data specfic hacks
-        label = process_labels_hacky(coords, shape, self.sigma, self.eps, bb=bb)
-        gap = 6
-        raw_patch, label = raw_patch[gap:-gap], label[gap:-gap]
+        # label = process_labels_hacky(coords, shape, self.sigma, self.eps, bb=bb)
+        # gap = 8
+        # raw_patch, label = raw_patch[gap:-gap], label[gap:-gap]
 
         have_label_channels = label.ndim == 4
         if have_label_channels: