Merge branch 'main' into more-inner-ear-analysis

Author: constantinpape

examples/domain_adaptation.py

@@ -0,0 +1,81 @@
+"""This script contains an example for using domain adptation to
+transfer a trained model for vesicle segmentation to a new dataset from a different data distribution,
+e.g. data from regular transmission electron microscopy (2D) instead of electron tomography or data from
+a different electron tomogram with different specimen and sample preparation.
+You don't need any annotations in the new domain to run this script.
+
+You can download example data for this script from:
+- Adaptation to 2d TEM data: TODO zenodo link
+- Adaptation to different tomography data: TODO zenodo link
+"""
+
+import os
+from glob import glob
+
+from sklearn.model_selection import train_test_split
+from synaptic_reconstruction.training import mean_teacher_adaptation
+from synaptic_reconstruction.tools.util import get_model_path
+
+
+def main():
+    # Choose whether to adapt the model to 2D or to 3D data.
+    train_2d_model = True
+
+    # TODO adjust to zenodo downloads
+    # These are the data folders for the example data downloaded from zenodo.
+    # Update these paths to apply the script to your own data.
+    # Check out the example data to see the data format for training.
+    data_root_folder_2d = "./data/2d_tem/train_unlabeled"
+    data_root_folder_3d = "./data/..."
+
+    # Choose the correct data folder depending on 2d/3d training.
+    data_root_folder = data_root_folder_2d if train_2d_model else data_root_folder_3d
+
+    # Get all files with ending .h5 in the training folder.
+    files = sorted(glob(os.path.join(data_root_folder, "**", "*.h5"), recursive=True))
+
+    # Crate a train / val split.
+    train_ratio = 0.85
+    train_paths, val_paths = train_test_split(files, test_size=1 - train_ratio, shuffle=True, random_state=42)
+
+    # Choose settings for the 2d or 3d domain adaptation.
+    if train_2d_model:
+        # This is the name of the checkpoint of the adapted model.
+        # For the name here the checkpoint will be stored in './checkpoints/example-2d-adapted-model'
+        model_name = "example-2d-adapted-model"
+        # The training patch size.
+        patch_shape = (256, 256)
+        # The batch size for training. You can increase this if you have enough VRAM.
+        batch_size = 4
+        # Get the checkpoint of the pretrained model for 2d vesicle segmentation.
+        source_checkpoint = get_model_path(model_type="vesicles_2d")
+    else:
+        # This is the name of the checkpoint of the adapted model.
+        # For the name here the checkpoint will be stored in './checkpoints/example-3d-adapted-model'
+        model_name = "example-3d-adapted-model"
+        # The training patch size.
+        patch_shape = (48, 256, 256)
+        # The batch size for training. You can increase this if you have enough VRAM.
+        batch_size = 1
+        # Get the checkpoint of the pretrained model for d vesicle segmentation.
+        source_checkpoint = get_model_path(model_type="vesicles_3d")
+
+    # We set the number of training iterations to 25,000.
+    n_iterations = int(2.5e4)
+
+    # This function runs the domain adaptation. Check out its documentation for
+    # advanced settings to update the training procedure.
+    mean_teacher_adaptation(
+        name=model_name,
+        unsupervised_train_paths=train_paths,
+        unsupervised_val_paths=val_paths,
+        source_checkpoint=source_checkpoint,
+        patch_shape=patch_shape,
+        batch_size=batch_size,
+        n_iterations=n_iterations,
+        confidence_threshold=0.75,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/network_training.py

@@ -0,0 +1,77 @@
+"""This script contains an example for how to train a network for
+a segmentation task with SynapseNet. This script covers the case of
+supervised training, i.e. your data needs to contain annotations for
+the structures you want to segment. If you want to use domain adaptation
+to adapt an already trained network to your data without the need for
+additional annotations then check out `domain_adaptation.py`.
+
+You can download example data for this script from:
+TODO zenodo link to Single-Ax / Chemical Fix data.
+"""
+import os
+from glob import glob
+
+from sklearn.model_selection import train_test_split
+from synaptic_reconstruction.training import supervised_training
+
+
+def main():
+    # This is the folder that contains your training data.
+    # The example was designed so that it runs for the sample data downloaded to './data'.
+    # If you want to train on your own data than change this filepath accordingly.
+    # TODO update to match zenodo download
+    data_root_folder = "./data/vesicles/train"
+
+    # The training data should be saved as .h5 files, with:
+    # an internal dataset called 'raw' that contains the image data
+    # and another dataset that contains the training annotations.
+    label_key = "labels/vesicles"
+
+    # Get all files with the ending .h5 in the training folder.
+    files = sorted(glob(os.path.join(data_root_folder, "**", "*.h5"), recursive=True))
+
+    # Crate a train / val split.
+    train_ratio = 0.85
+    train_paths, val_paths = train_test_split(files, test_size=1 - train_ratio, shuffle=True, random_state=42)
+
+    # We can either train a 2d or a 3d model. Whether a 2d or a 3d model is trained is derived from the patch shape.
+    # If your training data for 2d is stored as images (i.e. 2d data) them choose a  patch shape of form Y x X,
+    # e.g. (384, 384). If your data is stored in 3d, but you want to train a 2d model on it, choose a patch shape
+    # of the form 1 x Y x X, e.g. (1, 384, 384).
+    # If you want to train a 3d model then choose a patch shape of form Z x Y x X, e.g. (48, 256, 256).
+    train_2d_model = True
+    if train_2d_model:
+        batch_size = 2  # You can increase the batch size if you have enough VRAM.
+        # The model name determines the name of the checkpoint. E.g., for the name here the checkpoint will
+        # be saved at: 'checkpoints/example-2d-vesicle-model/'.
+        model_name = "example-2d-vesicle-model"
+        # The patch shape for training. See futher explanations above.
+        patch_shape = (1, 384, 384)
+    else:
+        batch_size = 1  # You can increase the batch size if you have enough VRAM.
+        # See the explanations for model_name and patch_shape above.
+        model_name = "example-3d-vesicle-model"
+        patch_shape = (48, 256, 256)
+
+    # If check_loader is set to True the training samples will be visualized via napari
+    # instead of starting a training. This is useful to validate that the training data
+    # is read correctly.
+    check_loader = False
+
+    # This function runs the training. Check out its documentation for
+    # advanced settings to update the training procedure.
+    supervised_training(
+        name=model_name,
+        train_paths=train_paths,
+        val_paths=val_paths,
+        label_key=label_key,
+        patch_shape=patch_shape,
+        batch_size=batch_size,
+        n_samples_train=None,
+        n_samples_val=25,
+        check=check_loader,
+    )
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -13,7 +13,7 @@
     license="MIT",
     entry_points={
         "console_scripts": [
-            "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli"
+            "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli",
         ],
         "napari.manifest": [
             "synaptic_reconstruction = synaptic_reconstruction:napari.yaml",

synaptic_reconstruction/tools/base_widget.py

@@ -61,7 +61,7 @@ def _update_selector(self, selector, layer_filter):
         image_layers = [layer.name for layer in self.viewer.layers if isinstance(layer, layer_filter)]  # if isinstance(layer, napari.layers.Image)
         selector.addItems(image_layers)
 
-    def _get_layer_selector_data(self, selector_name):
+    def _get_layer_selector_data(self, selector_name, return_metadata=False):
         """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]
@@ -72,7 +72,10 @@ def _get_layer_selector_data(self, selector_name):
             if isinstance(image_selector, QComboBox):
                 selected_layer_name = image_selector.currentText()
                 if selected_layer_name in self.viewer.layers:
-                    return self.viewer.layers[selected_layer_name].data
+                    if return_metadata:
+                        return self.viewer.layers[selected_layer_name].metadata
+                    else:
+                        return self.viewer.layers[selected_layer_name].data
         return None  # Return None if layer not found
 
     def _add_string_param(self, name, value, title=None, placeholder=None, layout=None, tooltip=None):
@@ -169,6 +172,15 @@ def _add_shape_param(self, names, values, min_val, max_val, step=1, title=None,
             title=title[1] if title is not None else title, tooltip=tooltip
         )
         layout.addLayout(y_layout)
+        
+        if len(names) == 3:
+            z_layout = QVBoxLayout()
+            z_param, _ = self._add_int_param(
+                names[2], values[2], min_val=min_val, max_val=max_val, layout=z_layout, step=step,
+                title=title[2] if title is not None else title, tooltip=tooltip
+            )
+            layout.addLayout(z_layout)
+            return x_param, y_param, z_param, layout
 
         return x_param, y_param, layout
 

synaptic_reconstruction/tools/cli.py

@@ -6,6 +6,7 @@
 
 
 # TODO: handle kwargs
+# TODO: add custom model path
 def segmentation_cli():
     parser = argparse.ArgumentParser(description="Run segmentation.")
     parser.add_argument(
@@ -16,6 +17,7 @@ def segmentation_cli():
         "--output_path", "-o", required=True,
         help="The filepath to directory where the segmentations will be saved."
     )
+    # TODO: list the availabel models here by parsing the keys of the model registry
     parser.add_argument(
         "--model", "-m", required=True, help="The model type."
     )

synaptic_reconstruction/tools/distance_measure_widget.py

@@ -101,33 +101,53 @@ def _add_lines_and_table(self, lines, properties, table_data, name):
     def on_measure_seg_to_object(self):
         segmentation = self._get_layer_selector_data(self.image_selector_name1)
         object_data = self._get_layer_selector_data(self.image_selector_name2)
+        # get metadata from layer if available
+        metadata = self._get_layer_selector_data(self.image_selector_name1, return_metadata=True)
+        resolution = metadata.get("voxel_size", None)
+        if resolution is not None:
+            resolution = [v for v in resolution.values()]
+        # if user input is present override metadata
+        if self.voxel_size_param.value() != 0.0:  # changed from default
+            resolution = segmentation.ndim * [self.voxel_size_param.value()]
 
         (distances,
          endpoints1,
          endpoints2,
          seg_ids) = distance_measurements.measure_segmentation_to_object_distances(
             segmentation=segmentation, segmented_object=object_data, distance_type="boundary",
+            resolution=resolution
         )
         lines, properties = distance_measurements.create_object_distance_lines(
             distances=distances,
             endpoints1=endpoints1,
             endpoints2=endpoints2,
-            seg_ids=seg_ids
+            seg_ids=seg_ids,
         )
         table_data = self._to_table_data(distances, seg_ids, endpoints1, endpoints2)
         self._add_lines_and_table(lines, properties, table_data, name="distances")
 
     def on_measure_pairwise(self):
         segmentation = self._get_layer_selector_data(self.image_selector_name1)
+        if segmentation is None:
+            show_info("Please choose a segmentation.")
+            return
+        # get metadata from layer if available
+        metadata = self._get_layer_selector_data(self.image_selector_name1, return_metadata=True)
+        resolution = metadata.get("voxel_size", None)
+        if resolution is not None:
+            resolution = [v for v in resolution.values()]
+        # if user input is present override metadata
+        if self.voxel_size_param.value() != 0.0:  # changed from default
+            resolution = segmentation.ndim * [self.voxel_size_param.value()]
 
         (distances,
          endpoints1,
          endpoints2,
          seg_ids) = distance_measurements.measure_pairwise_object_distances(
-            segmentation=segmentation, distance_type="boundary"
+            segmentation=segmentation, distance_type="boundary", resolution=resolution
         )
         lines, properties = distance_measurements.create_pairwise_distance_lines(
-            distances=distances, endpoints1=endpoints1, endpoints2=endpoints2, seg_ids=seg_ids.tolist()
+            distances=distances, endpoints1=endpoints1, endpoints2=endpoints2, seg_ids=seg_ids.tolist(),
         )
         table_data = self._to_table_data(
             distances=properties["distance"],
@@ -142,5 +162,10 @@ def _create_settings_widget(self):
         self.save_path, layout = self._add_path_param(name="Save Table", select_type="file", value="")
         setting_values.layout().addLayout(layout)
 
+        self.voxel_size_param, layout = self._add_float_param(
+            "voxel_size", 0.0, min_val=0.0, max_val=100.0,
+        )
+        setting_values.layout().addLayout(layout)
+
         settings = self._make_collapsible(widget=setting_values, title="Advanced Settings")
         return settings