Merge branch 'main' into add-mitov3

Author: lufre1

README.md

@@ -3,7 +3,7 @@
 SynapseNet is a tool for segmentation and analysis of synapses in electron microscopy.
 
 To learn how to use SynapseNet, check out [the documentation](https://computational-cell-analytics.github.io/synapse-net/) and [the tutorial video](https://youtu.be/7n8Oq1uAByE).
-To learn more about how it works, check out [our preprint](https://www.biorxiv.org/content/10.1101/2024.12.02.626387v1).
+To learn more about how it works, check out [our publication](https://www.molbiolcell.org/doi/full/10.1091/mbc.E24-11-0519).
 
 See an example reconstruction of a mossy fibre synapse with SynapseNet.
 Automatic segmentation of synaptic vesicles are rendered in orange, active zones in blue and two mitochondria in red and cyan.

synapse_net/inference/inference.py

@@ -114,6 +114,8 @@ def get_model_training_resolution(model_type: str) -> Dict[str, float]:
         "active_zone": {"x": 1.38, "y": 1.38, "z": 1.38},
         "compartments": {"x": 3.47, "y": 3.47, "z": 3.47},
         "mitochondria": {"x": 2.07, "y": 2.07, "z": 2.07},
+        # TODO: this is just copied from the previous mito model, it may be necessary to update this.
+        "mitochondria2": {"x": 1.45, "y": 1.45, "z": 1.45},
         "mitochondria3": {"x": 2.87, "y": 2.87, "z": 2.87},
         "cristae": {"x": 1.44, "y": 1.44, "z": 1.44},
         "ribbon": {"x": 1.188, "y": 1.188, "z": 1.188},

synapse_net/tools/cli.py

@@ -107,17 +107,23 @@ def pool_visualization_cli():
         help="The filepath to the mrc file containing the tomogram data."
     )
     parser.add_argument(
-        "--vesicle_path", "-v", required=True, help="The filepath to the tif file containing the vesicle segmentation."
+        "--vesicle_paths", "-v", required=True, nargs="+",
+        help="The filepath(s) to the tif file(s) containing the vesicle segmentation."
     )
     parser.add_argument(
         "--table_paths", "-t", required=True, nargs="+",
-        help="The filepath to the table with the vesicle pool assignments."
+        help="The filepath(s) to the table(s) with the vesicle pool assignments."
     )
     parser.add_argument(
         "-s", "--segmentation_paths", nargs="+", help="Filepaths for additional segmentations."
     )
+    parser.add_argument(
+        "--split_pools", action="store_true", help="Whether to split the pools into individual layers.",
+    )
     args = parser.parse_args()
-    _visualize_vesicle_pools(args.input_path, args.vesicle_path, args.table_paths, args.segmentation_paths)
+    _visualize_vesicle_pools(
+        args.input_path, args.vesicle_paths, args.table_paths, args.segmentation_paths, args.split_pools
+    )
 
 
 # TODO: handle kwargs

synapse_net/tools/pool_visualization.py

@@ -2,19 +2,34 @@
 
 import imageio.v3 as imageio
 import napari
+import numpy as np
 import pandas as pd
 
+from elf.parallel import isin
 from ..file_utils import read_mrc
 
 
-def _create_pools(vesicles, table):
+def _create_pools(vesicles, table, split_pools):
     label_ids, pool_colors = table.label.values, table.color.values
 
     pools = vesicles
     colormap = {label_id: color for label_id, color in zip(label_ids, pool_colors)}
     colormap[None] = [0, 0, 0, 0]
 
-    return pools, colormap
+    if split_pools:
+        unique_colors = np.unique(pool_colors)
+        pool_ret = {}
+        for this_color in unique_colors:
+            this_pool = pools.copy()
+            this_ids = [label_id for label_id, color in colormap.items() if color == this_color]
+            pool_mask = np.zeros(this_pool.shape, dtype="bool")
+            pool_mask = isin(this_pool, this_ids, out=pool_mask, block_shape=(32, 128, 128))
+            this_pool[~pool_mask] = 0
+            pool_ret[this_color] = this_pool
+    else:
+        pool_ret = {"pools": pools}
+
+    return pool_ret, colormap
 
 
 def _parse_tables(table_paths):
@@ -39,25 +54,36 @@ def load_table(path):
     return table
 
 
-def _visualize_vesicle_pools(input_path, vesicle_path, table_paths, segmentation_paths):
+def _visualize_vesicle_pools(input_path, vesicle_paths, table_paths, segmentation_paths, split_pools):
     # Load the tomogram data, including scale information.
     data, voxel_size = read_mrc(input_path)
     axes = "zyx" if data.ndim == 3 else "yx"
     scale = tuple(float(voxel_size[ax]) for ax in axes)
     print("Loading data with scale", scale, "nanometer")
 
-    # Load the vesicle layer.
-    vesicles = imageio.imread(vesicle_path)
+    # Load the vesicle layer, either from a single file with
+    if len(vesicle_paths) == 1:
+        vesicles = imageio.imread(vesicle_paths)
+    else:
+        vesicles = None
+        for path in vesicle_paths:
+            this_vesicles = imageio.imread(path)
+            if vesicles is None:
+                vesicles = this_vesicles.copy()
+            else:
+                ves_mask = this_vesicles != 0
+                vesicles[ves_mask] = this_vesicles[ves_mask]
 
     # Load the tables with the pool assignments.
     # Create and add the pool layer.
     table = _parse_tables(table_paths)
-    pools, colormap = _create_pools(vesicles, table)
+    pools, colormap = _create_pools(vesicles, table, split_pools)
 
     viewer = napari.Viewer()
     viewer.add_image(data, scale=scale)
     viewer.add_labels(vesicles, scale=scale)
-    viewer.add_labels(pools, scale=scale, name="pools", colormap=colormap)
+    for pool_name, pool in pools.items():
+        viewer.add_labels(pool, scale=scale, name=pool_name, colormap=colormap)
 
     # Add the additional segmentations.
     if segmentation_paths is not None:

test/test_inference.py

@@ -4,6 +4,7 @@
 from shutil import rmtree
 
 import imageio.v3 as imageio
+import requests
 from synapse_net.file_utils import read_mrc
 from synapse_net.sample_data import get_sample_data
 
@@ -23,6 +24,26 @@ def tearDown(self):
         except OSError:
             pass
 
+    # Test that all models can be accessed by pooch and have the necessary resolution information.
+    def test_models(self):
+        from synapse_net.inference.inference import _get_model_registry, get_model_training_resolution
+
+        def check_url(url):
+            try:
+                # Make a HEAD request to the URL, which fetches HTTP headers but no content.
+                response = requests.head(url, allow_redirects=True)
+                # Check if the HTTP status code is one that indicates availability (200 <= code < 400).
+                return response.status_code < 400
+            except requests.RequestException:
+                # Handle connection exceptions
+                return False
+
+        registry = _get_model_registry()
+        for name in registry.registry.keys():
+            url_exists = check_url(registry.get_url(name))
+            self.assertTrue(url_exists)
+            get_model_training_resolution(name)
+
     def test_run_segmentation(self):
         from synapse_net.inference import run_segmentation, get_model