Small API changes for training (#70)

C-Achard · web-flow · commit b569160ada6f · 2024-05-06T13:40:20.000+02:00
* Add optional train/val split for supervised models

* Update worker_training.py

* Change name to WNet3D

* Update worker_training.py

* Replace CellSeg3d with CellSeg3D

* Update test_training.py

* Update test_training.py

* Fix missing Wnet3d + fix tests (utils)
diff --git a/.napari/DESCRIPTION.md b/.napari/DESCRIPTION.md
@@ -117,7 +117,7 @@ this information here.
 ## Getting Help
 
 If you would like to report an issue with the plugin,
-please open an [issue on Github](https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues)
+please open an [issue on Github](https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues)
 <!--
 This section should point users to your preferred support tools, whether this be raising
 an issue on GitHub, asking a question on image.sc, or using some other method of contact.
diff --git a/.napari/config.yml b/.napari/config.yml
@@ -2,11 +2,11 @@ visibility: public
 # note : remember to add napari-hub preview app to repo
 
 project_urls:
-  Bug Tracker: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
-  Documentation: https://adaptivemotorcontrollab.github.io/CellSeg3d/res/welcome.html
-  Source Code: https://github.com/AdaptiveMotorControlLab/CellSeg3d
-  Project Site: https://github.com/AdaptiveMotorControlLab/CellSeg3d
-  User Support: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
+  Bug Tracker: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues
+  Documentation: https://adaptivemotorcontrollab.github.io/CellSeg3D/res/welcome.html
+  Source Code: https://github.com/AdaptiveMotorControlLab/CellSeg3D
+  Project Site: https://github.com/AdaptiveMotorControlLab/CellSeg3D
+  User Support: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues
 
 # Add labels from the EDAM Bioimaging ontology
 labels:
diff --git a/README.md b/README.md
@@ -5,9 +5,9 @@
 [![PyPI](https://img.shields.io/pypi/v/napari-cellseg3d.svg?color=green)](https://pypi.org/project/napari-cellseg3d)
 [![Downloads](https://static.pepy.tech/badge/napari-cellseg3d)](https://pepy.tech/project/napari-cellseg3d)
 [![Downloads](https://static.pepy.tech/badge/napari-cellseg3d/month)](https://pepy.tech/project/napari-cellseg3d)
-[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://github.com/AdaptiveMotorControlLab/CellSeg3d/raw/main/LICENSE)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://github.com/AdaptiveMotorControlLab/CellSeg3D/raw/main/LICENSE)
 [![Python Version](https://img.shields.io/pypi/pyversions/napari-cellseg-annotator.svg?color=green)](https://python.org)
-[![codecov](https://codecov.io/gh/AdaptiveMotorControlLab/CellSeg3d/branch/main/graph/badge.svg?token=hzUcn3XN8F)](https://codecov.io/gh/AdaptiveMotorControlLab/CellSeg3d)
+[![codecov](https://codecov.io/gh/AdaptiveMotorControlLab/CellSeg3D/branch/main/graph/badge.svg?token=hzUcn3XN8F)](https://codecov.io/gh/AdaptiveMotorControlLab/CellSeg3D)
 [![napari hub](https://img.shields.io/endpoint?url=https://api.napari-hub.org/shields/napari-cellseg3d)](https://www.napari-hub.org/plugins/napari-cellseg3d)
 
 - A napari plugin for 3D cell segmentation: training, inference, and data review. In particular, this project was developed for analysis of mesoSPIM-acquired (cleared tissue + lightsheet) datasets.
@@ -20,7 +20,7 @@
 
 ## Documentation
 
-📚 A lot of documentation is available at https://AdaptiveMotorControlLab.github.io/CellSeg3d
+📚 A lot of documentation is available at https://AdaptiveMotorControlLab.github.io/CellSeg3D
 
 You can also generate docs by running ``make html`` in the docs/ folder.
 
@@ -70,6 +70,7 @@ To avoid issues when installing on the ARM64 architecture, please follow these s
         conda env create -f conda/CellSeg3D_ARM64.yml
         conda activate napari_CellSeg3D_ARM64
 
+
 2) Install a Qt backend (PySide or PyQt5)
 3) Launch napari, the plugin should be available in the plugins menu.
 
@@ -137,7 +138,7 @@ Distributed under the terms of the [MIT] license.
 
 [napari-hub]: https://www.napari-hub.org/plugins/napari-cellseg3d
 
-[file an issue]: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
+[file an issue]: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues
 [napari]: https://github.com/napari/napari
 [Cookiecutter]: https://github.com/audreyr/cookiecutter
 [@napari]: https://github.com/napari
@@ -146,7 +147,7 @@ Distributed under the terms of the [MIT] license.
 [tox]: https://tox.readthedocs.io/en/latest/
 [pip]: https://pypi.org/project/pip/
 [PyPI]: https://pypi.org/
-[Installation page]: https://adaptivemotorcontrollab.github.io/CellSeg3d/source/guides/installation_guide.html
+[Installation page]: https://adaptivemotorcontrollab.github.io/CellSeg3D/source/guides/installation_guide.html
 [the PyTorch website for installation instructions]: https://pytorch.org/get-started/locally/
 [PyTorch]: https://pytorch.org/get-started/locally/
 [MONAI's optional dependencies]: https://docs.monai.io/en/stable/installation.html#installing-the-recommended-dependencies
diff --git a/docs/_config.yml b/docs/_config.yml
@@ -21,7 +21,7 @@ bibtex_bibfiles:
 
 # Information about where the book exists on the web
 repository:
-  url: https://github.com/AdaptiveMotorControlLab/CellSeg3d  # Online location of your book
+  url: https://github.com/AdaptiveMotorControlLab/CellSeg3D  # Online location of your book
   path_to_book: docs  # Optional path to your book, relative to the repository root
   branch: main  # Which branch of the repository should be used when creating links (optional)
 
diff --git a/docs/source/guides/custom_model_template.rst b/docs/source/guides/custom_model_template.rst
@@ -31,4 +31,4 @@ To add a custom model, you will need a **.py** file with the following structure
 .. note::
     **WIP** : Currently you must modify :doc:`model_framework.py <../code/_autosummary/napari_cellseg3d.code_models.model_framework>` as well : import your model class and add it to the ``model_dict`` attribute
 
-.. _file an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
+.. _file an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues
diff --git a/docs/source/guides/detailed_walkthrough.rst b/docs/source/guides/detailed_walkthrough.rst
@@ -207,7 +207,7 @@ To visualise some of the results when running on a folder, you can leave the **V
 .. note::
     Too plot your results, check out the `provided notebooks`_
 
-.. _provided notebooks: https://github.com/AdaptiveMotorControlLab/CellSeg3d/tree/main/notebooks
+.. _provided notebooks: https://github.com/AdaptiveMotorControlLab/CellSeg3D/tree/main/notebooks
 
 
 You can then launch inference and the results will be saved in your specified folder.
@@ -271,6 +271,6 @@ for the plots to work.
    Example of the plot present in the notebooks.
    Coordinates are based on centroids, the size represents the volume, the color, and the sphericity.
 
-.. _notebooks folder of the repository: https://github.com/AdaptiveMotorControlLab/CellSeg3d/tree/main/notebooks
+.. _notebooks folder of the repository: https://github.com/AdaptiveMotorControlLab/CellSeg3D/tree/main/notebooks
 
 With this complete, you can repeat the workflow as needed.
diff --git a/docs/source/guides/installation_guide.rst b/docs/source/guides/installation_guide.rst
@@ -8,7 +8,7 @@ Please refer to the :ref:`section below <source/guides/installation_guide:ARM64
 .. warning::
     If you encounter any issues during installation, feel free to open an issue on our `GitHub repository`_.
 
-.. _GitHub repository: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
+.. _GitHub repository: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues
 
 
 Installing pre-requisites
@@ -82,11 +82,11 @@ ______________________________
 
 .. _miniconda3: https://docs.conda.io/projects/conda/en/latest/user-guide/install/macos.html
 
-1. **Clone the repository** (`link <https://github.com/AdaptiveMotorControlLab/CellSeg3d>`_):
+1. **Clone the repository** (`link <https://github.com/AdaptiveMotorControlLab/CellSeg3D>`_):
 
 .. code-block::
 
-    git clone https://github.com/AdaptiveMotorControlLab/CellSeg3d.git
+    git clone https://github.com/AdaptiveMotorControlLab/CellSeg3D.git
 
 2. **Create the Conda Environment** :
 In the terminal, navigate to the CellSeg3D folder:
diff --git a/docs/source/guides/training_wnet.rst b/docs/source/guides/training_wnet.rst
@@ -96,4 +96,4 @@ Common issues troubleshooting
 
 
 .. _WNet, A Deep Model for Fully Unsupervised Image Segmentation: https://arxiv.org/abs/1711.08506
-.. _open an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
+.. _open an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues
diff --git a/examples/README.md b/examples/README.md
@@ -4,4 +4,5 @@ Here is a very small volume (from [IDR project 853](https://idr.openmicroscopy.o
 You can install, launch `napari`, activate the CellSeg3D plugin app, and drag & drop this volume into the canvas.
 Then, for example, run `inference` with one of our models.
 
-See [CellSeg3D documentation](https://adaptivemotorcontrollab.github.io/CellSeg3d/welcome.html) for more details.
+See [CellSeg3D documentation](https://adaptivemotorcontrollab.github.io/CellSeg3D/welcome.html) for more details.
+
diff --git a/napari_cellseg3d/_tests/test_plugin_inference.py b/napari_cellseg3d/_tests/test_plugin_inference.py
@@ -31,7 +31,7 @@ def test_inference(make_napari_viewer_proxy, qtbot):
 
     assert widget.check_ready()
 
-    widget.model_choice.setCurrentText("WNet")
+    widget.model_choice.setCurrentText("WNet3D")
     widget._restrict_window_size_for_model()
     assert widget.use_window_choice.isChecked()
     assert widget.window_size_choice.currentText() == "64"
diff --git a/napari_cellseg3d/_tests/test_plugin_training.py b/napari_cellseg3d/_tests/test_plugin_training.py
@@ -33,7 +33,7 @@ def test_worker_configs(make_napari_viewer_proxy):
                 worker.config, attr
             )
     # test unsupervised config and worker
-    widget.model_choice.setCurrentText("WNet")
+    widget.model_choice.setCurrentText("WNet3D")
     widget._toggle_unsupervised_mode(enabled=True)
     default_config = config.WNetTrainingWorkerConfig()
     worker = widget._create_worker(additional_results_description="TEST_1")
diff --git a/napari_cellseg3d/_tests/test_plugin_utils.py b/napari_cellseg3d/_tests/test_plugin_utils.py
@@ -38,7 +38,7 @@ def test_crop_widget(make_napari_viewer_proxy):
     view = make_napari_viewer_proxy()
     widget = Cropping(view)
 
-    image = rand_gen.random((10, 10, 10)).astype(np.uint8)
+    image = rand_gen.random((10, 10, 10)).astype(np.int8)
     image_layer_1 = view.add_image(image, name="image")
     image_layer_2 = view.add_labels(image, name="image2")
 
diff --git a/napari_cellseg3d/_tests/test_training.py b/napari_cellseg3d/_tests/test_training.py
@@ -86,8 +86,8 @@ def test_unsupervised_training(make_napari_viewer_proxy):
     widget = Trainer(viewer)
     widget.log = LogFixture()
     widget.worker = None
+    widget.model_choice.setCurrentText("WNet3D")
     widget._toggle_unsupervised_mode(enabled=True)
-    widget.model_choice.setCurrentText("WNet")
 
     widget.patch_choice.setChecked(True)
     [w.setValue(4) for w in widget.patch_size_widgets]
diff --git a/napari_cellseg3d/code_models/models/pretrained/pretrained_model_urls.json b/napari_cellseg3d/code_models/models/pretrained/pretrained_model_urls.json
@@ -4,6 +4,7 @@
   "VNet": "https://huggingface.co/C-Achard/cellseg3d/resolve/main/VNet_latest.tar.gz",
   "SwinUNetR": "https://huggingface.co/C-Achard/cellseg3d/resolve/main/SwinUNetR_latest.tar.gz",
   "WNet": "https://huggingface.co/C-Achard/cellseg3d/resolve/main/wnet_latest.tar.gz",
+  "WNet3D": "https://huggingface.co/C-Achard/cellseg3d/resolve/main/wnet_latest.tar.gz",
   "WNet_ONNX": "https://huggingface.co/C-Achard/cellseg3d/resolve/main/wnet_onnx.tar.gz",
   "test": "https://huggingface.co/C-Achard/cellseg3d/resolve/main/test.tar.gz"
 }
diff --git a/napari_cellseg3d/code_models/models/wnet/model.py b/napari_cellseg3d/code_models/models/wnet/model.py
@@ -17,7 +17,7 @@
 
 
 class WNet_encoder(nn.Module):
-    """WNet with encoder only."""
+    """WNet3D with encoder only."""
 
     def __init__(
         self,
diff --git a/napari_cellseg3d/code_models/worker_inference.py b/napari_cellseg3d/code_models/worker_inference.py
@@ -277,7 +277,7 @@ def load_layer(self):
 
         normalization = (
             QuantileNormalization()
-            if self.config.model_info.name != "WNet"
+            if self.config.model_info.name != "WNet3D"
             else lambda x: x
         )
         volume = np.reshape(volume, newshape=(1, *volume.shape))
diff --git a/napari_cellseg3d/code_models/worker_training.py b/napari_cellseg3d/code_models/worker_training.py
@@ -431,9 +431,9 @@ def train(
                 wandb.init(
                     config=config_dict,
                     project="CellSeg3D - WNet",
-                    name=f"WNet_training - {utils.get_date_time()}",
+                    name=f"WNet3D_training - {utils.get_date_time()}",
                     mode=self.wandb_config.mode,
-                    tags=["WNet", "training"],
+                    tags=["WNet3D", "training"],
                 )
 
             set_determinism(seed=self.config.deterministic_config.seed)
@@ -487,7 +487,7 @@ def train(
                     )
 
                     weights_file = WNet_.weights_file
-                    self.downloader.download_weights("WNet", weights_file)
+                    self.downloader.download_weights("WNet3D", weights_file)
                     weights = str(PRETRAINED_WEIGHTS_DIR / Path(weights_file))
                     self.config.weights_info.path = weights
 
@@ -781,7 +781,7 @@ def train(
 
             if WANDB_INSTALLED and self.wandb_config.save_model_artifact:
                 model_artifact = wandb.Artifact(
-                    "WNet",
+                    "WNet3D",
                     type="model",
                     description="CellSeg3D WNet",
                     metadata=self.config.__dict__,
@@ -1127,6 +1127,11 @@ def train(
         weights_config = self.config.weights_info
         deterministic_config = self.config.deterministic_config
 
+        if self.config.device == "mps":
+            from os import environ
+
+            environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
         start_time = time.time()
 
         try:
diff --git a/napari_cellseg3d/code_plugins/plugin_helper.py b/napari_cellseg3d/code_plugins/plugin_helper.py
@@ -21,10 +21,10 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         """Creates a widget with links to documentation and about page."""
         super().__init__()
 
-        self.help_url = "https://adaptivemotorcontrollab.github.io/CellSeg3d/"
+        self.help_url = "https://adaptivemotorcontrollab.github.io/CellSeg3D/"
 
         self.about_url = "https://wysscenter.ch/advances/3d-computer-vision-for-brain-analysis"
-        self.repo_url = "https://github.com/AdaptiveMotorControlLab/CellSeg3d"
+        self.repo_url = "https://github.com/AdaptiveMotorControlLab/CellSeg3D"
         self._viewer = viewer
 
         logo_path = str(
diff --git a/napari_cellseg3d/code_plugins/plugin_model_inference.py b/napari_cellseg3d/code_plugins/plugin_model_inference.py
@@ -258,7 +258,7 @@ def __init__(self, viewer: "napari.viewer.Viewer", parent=None):
         self._remove_unused()
 
     def _toggle_crf_choice(self):
-        if self.model_choice.currentText() == "WNet":
+        if self.model_choice.currentText() == "WNet3D":
             self.use_crf.setVisible(True)
         else:
             self.use_crf.setVisible(False)
@@ -335,7 +335,7 @@ def check_ready(self):
     def _restrict_window_size_for_model(self):
         """Sets the window size to a value that is compatible with the chosen model."""
         self.wnet_enabled = False
-        if self.model_choice.currentText() == "WNet":
+        if self.model_choice.currentText() == "WNet3D":
             self.wnet_enabled = True
             self.window_size_choice.setCurrentIndex(self._default_window_size)
             self.use_window_choice.setChecked(self.wnet_enabled)
diff --git a/napari_cellseg3d/code_plugins/plugin_model_training.py b/napari_cellseg3d/code_plugins/plugin_model_training.py
@@ -451,7 +451,7 @@ def check_ready(self):
 
     def _toggle_unsupervised_mode(self, enabled=False):
         """Change all the UI elements needed for unsupervised learning mode."""
-        if self.model_choice.currentText() == "WNet" or enabled:
+        if self.model_choice.currentText() == "WNet3D" or enabled:
             unsupervised = True
             self.start_btn = self.start_button_unsupervised
             if self.image_filewidget.text_field.text() == "Images directory":
@@ -799,7 +799,7 @@ def _build(self):
         ui.add_blank(advanced_tab, advanced_tab.layout)
         ##################
         model_params_group_w, model_params_group_l = ui.make_group(
-            "WNet parameters", r=20, b=5, t=11
+            "WNet3D parameters", r=20, b=5, t=11
         )
         ui.add_widgets(
             model_params_group_l,
diff --git a/napari_cellseg3d/config.py b/napari_cellseg3d/config.py
@@ -26,7 +26,7 @@
     "VNet": VNet_,
     "TRAILMAP_MS": TRAILMAP_MS_,
     "SwinUNetR": SwinUNETR_,
-    "WNet": WNet_,
+    "WNet3D": WNet_,
     # "TRAILMAP": TRAILMAP,
     # "test" : DO NOT USE, reserved for testing
 }
diff --git a/napari_cellseg3d/dev_scripts/classifier_test.ipynb b/napari_cellseg3d/dev_scripts/classifier_test.ipynb
diff --git a/napari_cellseg3d/dev_scripts/colab_training.py b/napari_cellseg3d/dev_scripts/colab_training.py
diff --git a/napari_cellseg3d/dev_scripts/test_new_evaluation.ipynb b/napari_cellseg3d/dev_scripts/test_new_evaluation.ipynb
diff --git a/notebooks/colab_wnet_training.ipynb b/notebooks/colab_wnet_training.ipynb
diff --git a/pyproject.toml b/pyproject.toml

Original file line number	Diff line number	Diff line change
`@@ -96,4 +96,4 @@ Common issues troubleshooting`
`96`	`96`
`97`	`97`
`98`	`98`	`.. _WNet, A Deep Model for Fully Unsupervised Image Segmentation: https://arxiv.org/abs/1711.08506`
`99`		`-.. _open an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues`
	`99`	`+.. _open an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3D/issues`