Merge pull request #51 from computational-cell-analytics/more-plugin-updates

Updates to napari plugin and start to implement CLI

Author: constantinpape

run_correction.sh

@@ -8,13 +8,12 @@
     name="synaptic_reconstruction",
     packages=find_packages(exclude=["test"]),
     version=__version__,
-    author="Constantin Pape; Sarah Muth",
+    author="Constantin Pape; Sarah Muth; Luca Freckmann",
     url="https://github.com/computational-cell-analytics/synaptic_reconstruction",
     license="MIT",
     entry_points={
         "console_scripts": [
-            "sr_tools.correct_segmentation = synaptic_reconstruction.tools.segmentation_correction:main",
-            "sr_tools.measure_distances = synaptic_reconstruction.tools.distance_measurement:main",
+            "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli"
         ],
         "napari.manifest": [
             "synaptic_reconstruction = synaptic_reconstruction:napari.yaml",

setup.py

@@ -1,6 +1,6 @@
 import os
 import multiprocessing as mp
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -61,8 +61,9 @@ def _compute_boundary_distances(segmentation, resolution, n_threads):
     n = len(seg_ids)
 
     pairwise_distances = np.zeros((n, n))
-    end_points1 = np.zeros((n, n, 3), dtype="int")
-    end_points2 = np.zeros((n, n, 3), dtype="int")
+    ndim = segmentation.ndim
+    end_points1 = np.zeros((n, n, ndim), dtype="int")
+    end_points2 = np.zeros((n, n, ndim), dtype="int")
 
     properties = regionprops(segmentation)
     properties = {prop.label: prop for prop in properties}
@@ -80,8 +81,11 @@ def compute_distances_for_object(i):
             prop = properties[ngb_id]
 
             bb = prop.bbox
-            offset = np.array(bb[:3])
-            bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
+            offset = np.array(bb[:ndim])
+            if ndim == 2:
+                bb = np.s_[bb[0]:bb[2], bb[1]:bb[3]]
+            else:
+                bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
 
             mask = segmentation[bb] == ngb_id
             ngb_dist, ngb_index = distances[bb].copy(), indices[(slice(None),) + bb]
@@ -168,10 +172,9 @@ def _compute_seg_object_distances(segmentation, segmented_object, resolution, ve
     for prop in tqdm(props, disable=not verbose):
         bb = prop.bbox
         offset = np.array(bb[:ndim])
-        # bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
-        if len(bb) == 4:  # 2D bounding box
+        if ndim == 2:
             bb = np.s_[bb[0]:bb[2], bb[1]:bb[3]]
-        elif len(bb) == 6:  # 3D bounding box
+        else:
             bb = np.s_[bb[0]:bb[3], bb[1]:bb[4], bb[2]:bb[5]]
 
         label = prop.label
@@ -296,7 +299,7 @@ def create_pairwise_distance_lines(
     distances: np.ndarray,
     endpoints1: np.ndarray,
     endpoints2: np.ndarray,
-    seg_ids: np.ndarray,
+    seg_ids: List[List[int]],
     n_neighbors: Optional[int] = None,
     pairs: Optional[np.ndarray] = None,
     bb: Optional[Tuple[slice]] = None,
@@ -323,9 +326,7 @@ def create_pairwise_distance_lines(
     if pairs is None and n_neighbors is not None:
         pairs = _extract_nearest_neighbors(distances, seg_ids, n_neighbors, remove_duplicates=remove_duplicates)
     elif pairs is None:
-        pairs = [
-            [id1, id2] for id1 in seg_ids for id2 in seg_ids if id1 < id2
-        ]
+        pairs = [[id1, id2] for id1 in seg_ids for id2 in seg_ids if id1 < id2]
 
     assert pairs is not None
     pair_indices = (

synaptic_reconstruction/distance_measurements.py

@@ -14,30 +14,43 @@ def _run_segmentation(
     block_shape=(128, 256, 256),
     halo=(48, 48, 48)
 ):
-
-    # get the segmentation via seeded watershed
     t0 = time.time()
-    seeds = parallel.label((foreground - boundaries) > 0.5, block_shape=block_shape, verbose=verbose)
+    boundary_threshold = 0.25
+    dist = parallel.distance_transform(
+        boundaries < boundary_threshold, halo=halo, verbose=verbose, block_shape=block_shape
+    )
     if verbose:
-        print("Compute connected components in", time.time() - t0, "s")
+        print("Compute distance transform in", time.time() - t0, "s")
 
+    # Get the segmentation via seeded watershed.
     t0 = time.time()
-    dist = parallel.distance_transform(seeds == 0, halo=halo, verbose=verbose, block_shape=block_shape)
+    seed_distance = 6
+    seeds = np.logical_and(foreground > 0.5, dist > seed_distance)
+    seeds = parallel.label(seeds, block_shape=block_shape, verbose=verbose)
     if verbose:
-        print("Compute distance transform in", time.time() - t0, "s")
+        print("Compute connected components in", time.time() - t0, "s")
+
+    # import napari
+    # v = napari.Viewer()
+    # v.add_image(boundaries)
+    # v.add_image(dist)
+    # v.add_labels(seeds)
+    # napari.run()
 
     t0 = time.time()
+    hmap = boundaries + ((dist.max() - dist) / dist.max())
     mask = (foreground + boundaries) > 0.5
+
     seg = np.zeros_like(seeds)
     seg = parallel.seeded_watershed(
-        dist, seeds, block_shape=block_shape,
+        hmap, seeds, block_shape=block_shape,
         out=seg, mask=mask, verbose=verbose, halo=halo,
     )
     if verbose:
         print("Compute watershed in", time.time() - t0, "s")
 
     seg = apply_size_filter(seg, min_size, verbose=verbose, block_shape=block_shape)
-    seg = _postprocess_seg_3d(seg)
+    seg = _postprocess_seg_3d(seg, area_threshold=5000)
     return seg
 
 

synaptic_reconstruction/inference/mitochondria.py

@@ -255,21 +255,23 @@ def _get_file_paths(input_path, ext=".mrc"):
 
 
 def _load_input(img_path, extra_files, i):
-    # Load the input data data
-    with open_file(img_path, "r") as f:
-
-        # Try to automatically derive the key with the raw data.
-        keys = list(f.keys())
-        if len(keys) == 1:
-            key = keys[0]
-        elif "data" in keys:
-            key = "data"
-        elif "raw" in keys:
-            key = "raw"
-
-        input_volume = f[key][:]
-    assert input_volume.ndim == 3
+    # Load the input data.
+    if os.path.splitext(img_path)[-1] == ".tif":
+        input_volume = imageio.imread(img_path)
 
+    else:
+        with open_file(img_path, "r") as f:
+            # Try to automatically derive the key with the raw data.
+            keys = list(f.keys())
+            if len(keys) == 1:
+                key = keys[0]
+            elif "data" in keys:
+                key = "data"
+            elif "raw" in keys:
+                key = "raw"
+            input_volume = f[key][:]
+
+    assert input_volume.ndim in (2, 3)
     # For now we assume this is always tif.
     if extra_files is not None:
         extra_input = imageio.imread(extra_files[i])
@@ -470,7 +472,7 @@ def apply_size_filter(
     return segmentation
 
 
-def _postprocess_seg_3d(seg):
+def _postprocess_seg_3d(seg, area_threshold=1000, iterations=4, iterations_3d=8):
     # Structure lement for 2d dilation in 3d.
     structure_element = np.ones((3, 3))  # 3x3 structure for XY plane
     structure_3d = np.zeros((1, 3, 3))  # Only applied in the XY plane
@@ -483,9 +485,9 @@ def _postprocess_seg_3d(seg):
         mask = seg[bb] == prop.label
 
         # Fill small holes and apply closing.
-        mask = remove_small_holes(mask, area_threshold=1000)
-        mask = np.logical_or(binary_closing(mask, iterations=4), mask)
-        mask = np.logical_or(binary_closing(mask, iterations=8, structure=structure_3d), mask)
+        mask = remove_small_holes(mask, area_threshold=area_threshold)
+        mask = np.logical_or(binary_closing(mask, iterations=iterations), mask)
+        mask = np.logical_or(binary_closing(mask, iterations=iterations_3d, structure=structure_3d), mask)
         seg[bb][mask] = prop.label
 
     return seg

synaptic_reconstruction/inference/util.py

@@ -5,10 +5,10 @@ categories: ["Image Processing", "Annotation"]
 contributions:
   commands:
     - id: synaptic_reconstruction.segment
-      python_name: synaptic_reconstruction.tools.synaptic_plugin.segmentation_widget:get_segmentation_widget
+      python_name: synaptic_reconstruction.tools.segmentation_widget:SegmentationWidget
       title: Segment
     - id: synaptic_reconstruction.distance_measure
-      python_name: synaptic_reconstruction.tools.synaptic_plugin.distance_measure_widget:get_distance_measure_widget
+      python_name: synaptic_reconstruction.tools.distance_measure_widget:DistanceMeasureWidget
       title: Distance Measurement
     - id: synaptic_reconstruction.file_reader
       title: Read volumetric data

synaptic_reconstruction/napari.yaml

@@ -1,9 +1,12 @@
 from pathlib import Path
+
 import napari
-from qtpy.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QLabel, QSpinBox, QLineEdit, QGroupBox, QFormLayout, QFrame, QComboBox, QCheckBox
 import qtpy.QtWidgets as QtWidgets
+
+from qtpy.QtWidgets import (
+    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QSpinBox, QComboBox, QCheckBox
+)
 from superqt import QCollapsible
-from magicgui.widgets import create_widget
 
 
 class BaseWidget(QWidget):
@@ -15,7 +18,7 @@ def __init__(self):
     def _create_layer_selector(self, selector_name, layer_type="Image"):
         """
         Create a layer selector for an image or labels and store it in a dictionary.
-        
+
         Parameters:
         - selector_name (str): The name of the selector, used as a key in the dictionary.
         - layer_type (str): The type of layer to filter for ("Image" or "Labels").
@@ -34,24 +37,24 @@ def _create_layer_selector(self, selector_name, layer_type="Image"):
         selector_widget = QtWidgets.QWidget()
         image_selector = QtWidgets.QComboBox()
         layer_label = QtWidgets.QLabel(f"{selector_name} Layer:")
-        
+
         # Populate initial options
         self._update_selector(selector=image_selector, layer_filter=layer_filter)
-        
+
         # Update selector on layer events
         self.viewer.layers.events.inserted.connect(lambda event: self._update_selector(image_selector, layer_filter))
         self.viewer.layers.events.removed.connect(lambda event: self._update_selector(image_selector, layer_filter))
 
         # Store the selector in the dictionary
         self.layer_selectors[selector_name] = selector_widget
-        
+
         # Set up layout
         layout = QVBoxLayout()
         layout.addWidget(layer_label)
         layout.addWidget(image_selector)
         selector_widget.setLayout(layout)
         return selector_widget
-    
+
     def _update_selector(self, selector, layer_filter):
         """Update a single selector with the current image layers in the viewer."""
         selector.clear()
@@ -62,10 +65,10 @@ def _get_layer_selector_data(self, selector_name):
         """Return the data for the layer currently selected in a given selector."""
         if selector_name in self.layer_selectors:
             selector_widget = self.layer_selectors[selector_name]
-            
+
             # Retrieve the QComboBox from the QWidget's layout
             image_selector = selector_widget.layout().itemAt(1).widget()
-            
+
             if isinstance(image_selector, QComboBox):
                 selected_layer_name = image_selector.currentText()
                 if selected_layer_name in self.viewer.layers:
@@ -176,7 +179,7 @@ def _make_collapsible(self, widget, title):
         collapsible.addWidget(widget)
         parent_widget.layout().addWidget(collapsible)
         return parent_widget
-    
+
     def _add_boolean_param(self, name, value, title=None, tooltip=None):
         checkbox = QCheckBox(name if title is None else title)
         checkbox.setChecked(value)

synaptic_reconstruction/tools/synaptic_plugin/base_widget.py renamed to synaptic_reconstruction/tools/base_widget.py

@@ -0,0 +1,53 @@
+import argparse
+from functools import partial
+
+from .util import run_segmentation, get_model
+from ..inference.util import inference_helper, parse_tiling
+
+
+# TODO: handle kwargs
+def segmentation_cli():
+    parser = argparse.ArgumentParser(description="Run segmentation.")
+    parser.add_argument(
+        "--input_path", "-i", required=True,
+        help="The filepath to the mrc file or the directory containing the tomogram data."
+    )
+    parser.add_argument(
+        "--output_path", "-o", required=True,
+        help="The filepath to directory where the segmentations will be saved."
+    )
+    parser.add_argument(
+        "--model", "-m", required=True, help="The model type."
+    )
+    parser.add_argument(
+        "--mask_path", help="The filepath to a tif file with a mask that will be used to restrict the segmentation."
+        "Can also be a directory with tifs if the filestructure matches input_path."
+    )
+    parser.add_argument("--input_key", "-k", required=False)
+    parser.add_argument(
+        "--force", action="store_true",
+        help="Whether to over-write already present segmentation results."
+    )
+    parser.add_argument(
+        "--tile_shape", type=int, nargs=3,
+        help="The tile shape for prediction. Lower the tile shape if GPU memory is insufficient."
+    )
+    parser.add_argument(
+        "--halo", type=int, nargs=3,
+        help="The halo for prediction. Increase the halo to minimize boundary artifacts."
+    )
+    parser.add_argument(
+        "--data_ext", default=".mrc", help="The extension of the tomogram data. By default .mrc."
+    )
+    args = parser.parse_args()
+
+    model = get_model(args.model)
+    tiling = parse_tiling(args.tile_shape, args.halo)
+
+    segmentation_function = partial(
+        run_segmentation, model=model, model_type=args.model, verbose=False, tiling=tiling,
+    )
+    inference_helper(
+        args.input_path, args.output_path, segmentation_function,
+        mask_input_path=args.mask_path, force=args.force, data_ext=args.data_ext,
+    )