WNet + models code refactor (#36)

* Testing instance methods Co-Authored-By: gityves <[email protected]> * Many fixes - Fixed monai reqs - Added custom functions for label checking - Fixed return type of voronoi_otsu and utils.resize - black * black * Complete instance method evaluation * Enfore pre-commit style * Removing dask-image * Fixed erroneous dtype conversion * Update test_plugin_utils.py * Update tox.ini * Added new pre-commit hooks * Run full suite of pre-commit hooks * Enforce style * Documentation update, crop contrast fix * Updated hooks * Update setup.cfg * Many fixes - Fixed monai reqs - Added custom functions for label checking - Fixed return type of voronoi_otsu and utils.resize - black * Enfore pre-commit style * Updated project files * Removing dask-image * Latest pre-commit hooks * Instance segmentation refactor + Voronoi-Otsu - Improved code for instance segmentation - Added Voronoi-Otsu labeling from pyclesperanto TODO : credits for labeling * isort * Fix inference * Added labeling tools + UI tweaks - Added tools from MLCourse to evaluate labels and auto-correct them - Instance seg benchmark notebook - Tweaked utils UI to scale according to Viewer size Co-Authored-By: gityves <[email protected]> * Many fixes - Fixed monai reqs - Added custom functions for label checking - Fixed return type of voronoi_otsu and utils.resize - black * Added pre-commit hooks * Update .pre-commit-config.yaml * Update pyproject.toml * Update pyproject.toml Ruff config * Enfore pre-commit style * Update .gitignore * Version bump * Revert "Version bump" This reverts commit 6e39971. * Updated project files * Fixed wrong value in instance sliders * Removing dask-image * Update test_plugin_utils.py * Relabeling tests * Added new pre-commit hooks * Latest pre-commit hooks * Run full suite of pre-commit hooks * Model class refactor * Added LR scheduler in training - Added ReduceLROnPlateau with params in training - Updated training guide - Minor UI attribute refactor - black * Update assess_instance.ipynb * Update .gitignore * Started adding WNet * Specify no grad in inference * First functional WNet inference, no CRF * Create test_models.py * Run full suite of pre-commit hooks * Patch for tests action + style * Add softNCuts basic test * Added crf Co-Authored-By: Nevexios <[email protected]> * More pre-commit checks * Functional CRF * Fix erroneous test comment, added toggle for crf - Warn if crf not installed - Fix test * Specify missing test deps * Trying to fix deps on Git * Removed master link to pydensecrf * Use commit hash * Removed commit hash * Removed master * Update tox.ini * Update pyproject.toml * Fixes and improvements - More CRF info - Added error handling to scheduler rate - Added ETA to training - Updated padding warning trigger size * Fixes and channel labeling prototype * Fixes - Fixed multi-channel instance and csv stats - Fixed rotation of inference outputs - Raised max crop size * Update plugin_model_inference.py * Update plugin_crop.py * Fixed patch_func sample number mismatch * Testing relabel tools * Fixes in inference * add model template + fix test + wnet loading opti - test fixes - changed crf input reqs - adapted instance seg for several channels * Update model_WNet.py * Update model_VNet.py * Fixed folder creation when saving to folder * Fix check_ready for results filewidget * Added remapping in WNet + ruff config * Run new hooks * Small docs update * Testing fix * Fixed multithread testing (locally) * Added proper tests for train/infer * Slight coverage increase * Update test_plugin_inference.py * Set window inference to 64 for WNet * Update instance_segmentation.py * Moved normalization to the correct place * Added auto-set dims for cropping * Update test_plugin_utils.py * More WNet - Added experimental .pt loading for jit models - More CRF tests - Optimized WNet by loading inference only * Update crf test/deps for testing * Update test_and_deploy.yml * Update test_and_deploy.yml * Update tox.ini * Update test_and_deploy.yml * Trying to fix tox install of pydensecrf * Added experimental ONNX support for inference * Updated WNet for ONNX conversion * Added dropout param * Minor fixes in training * Fix weights file extension in inference + coverage - Remove unused scripts - More tests - Fixed weights type in inference * Run all hooks * Fix inference testing * Changed anisotropy calculation * Finish rebase + bump version * Instance segmentation refactor + Voronoi-Otsu - Improved code for instance segmentation - Added Voronoi-Otsu labeling from pyclesperanto TODO : credits for labeling * Disabled small removal in Voronoi-Otsu * Added new docs for instance seg * Docs + UI update - Updated welcome/README - Changed step for DoubleCounter * Update requirements.txt Fix typo * isort * Fix tests * Fixed parental issues and instance seg widget init - Fixed widgets parents that were incorrectly init - Improve use of instance seg. method classes and init * Fix inference * Added labeling tools + UI tweaks - Added tools from MLCourse to evaluate labels and auto-correct them - Instance seg benchmark notebook - Tweaked utils UI to scale according to Viewer size Co-Authored-By: gityves <[email protected]> * Testing instance methods Co-Authored-By: gityves <[email protected]> * Many fixes - Fixed monai reqs - Added custom functions for label checking - Fixed return type of voronoi_otsu and utils.resize - black * black * Complete instance method evaluation * Added pre-commit hooks * Enfore pre-commit style * Update .gitignore * Version bump * Updated project files * Fixed missing parent error * Fixed wrong value in instance sliders * Removing dask-image * Fixed erroneous dtype conversion * Update test_plugin_utils.py * Temporary test action patch * Update plugin_convert.py * Update tox.ini Added pocl for testing on GH Actions * Update tox.ini * Found existing pocl * Updated utils test to avoid Voronoi-Otsu VO is missing CL runtime * Relabeling tests * Latest pre-commit hooks * Run full suite of pre-commit hooks * Enforce style * Instance segmentation refactor + Voronoi-Otsu - Improved code for instance segmentation - Added Voronoi-Otsu labeling from pyclesperanto TODO : credits for labeling * Disabled small removal in Voronoi-Otsu * Added new docs for instance seg * Docs + UI update - Updated welcome/README - Changed step for DoubleCounter * Update requirements.txt Fix typo * isort * Fix tests * Fixed parental issues and instance seg widget init - Fixed widgets parents that were incorrectly init - Improve use of instance seg. method classes and init * Fix inference * Added labeling tools + UI tweaks - Added tools from MLCourse to evaluate labels and auto-correct them - Instance seg benchmark notebook - Tweaked utils UI to scale according to Viewer size Co-Authored-By: gityves <[email protected]> * Testing instance methods Co-Authored-By: gityves <[email protected]> * Many fixes - Fixed monai reqs - Added custom functions for label checking - Fixed return type of voronoi_otsu and utils.resize - black * black * Complete instance method evaluation * Added pre-commit hooks * Update .pre-commit-config.yaml * Update pyproject.toml * Update pyproject.toml Ruff config * Enfore pre-commit style * Update .gitignore * Version bump * Revert "Version bump" This reverts commit 6e39971. * Updated project files * Fixed missing parent error * Fixed wrong value in instance sliders * Removing dask-image * Fixed erroneous dtype conversion * Update test_plugin_utils.py * Update plugin_convert.py * Update tox.ini Added pocl for testing on GH Actions * Update tox.ini * Found existing pocl * Updated utils test to avoid Voronoi-Otsu VO is missing CL runtime * Relabeling tests * Added new pre-commit hooks * Latest pre-commit hooks * Run full suite of pre-commit hooks * Model class refactor * Added LR scheduler in training - Added ReduceLROnPlateau with params in training - Updated training guide - Minor UI attribute refactor - black * Update assess_instance.ipynb * Update .gitignore * Started adding WNet * Specify no grad in inference * First functional WNet inference, no CRF * Create test_models.py * Run full suite of pre-commit hooks * Patch for tests action + style * Add softNCuts basic test * Added crf Co-Authored-By: Nevexios <[email protected]> * More pre-commit checks * Functional CRF * Fix erroneous test comment, added toggle for crf - Warn if crf not installed - Fix test * Specify missing test deps * Trying to fix deps on Git * Removed master link to pydensecrf * Use commit hash * Removed commit hash * Removed master * Update tox.ini * Update pyproject.toml * Fixes and improvements - More CRF info - Added error handling to scheduler rate - Added ETA to training - Updated padding warning trigger size * Fixes and channel labeling prototype * Fixes - Fixed multi-channel instance and csv stats - Fixed rotation of inference outputs - Raised max crop size * Update plugin_model_inference.py * Fixed patch_func sample number mismatch * Testing relabel tools * Fixes in inference * add model template + fix test + wnet loading opti - test fixes - changed crf input reqs - adapted instance seg for several channels * Update model_WNet.py * Update model_VNet.py * Fixed folder creation when saving to folder * Fix check_ready for results filewidget * Added remapping in WNet + ruff config * Run new hooks * Small docs update * Testing fix * Fixed multithread testing (locally) * Added proper tests for train/infer * Slight coverage increase * Update test_plugin_inference.py * Set window inference to 64 for WNet * Moved normalization to the correct place * Added auto-set dims for cropping * Update test_plugin_utils.py * More WNet - Added experimental .pt loading for jit models - More CRF tests - Optimized WNet by loading inference only * Update crf test/deps for testing * Update test_and_deploy.yml * Update test_and_deploy.yml * Update tox.ini * Update test_and_deploy.yml * Trying to fix tox install of pydensecrf * Added experimental ONNX support for inference * Updated WNet for ONNX conversion * Added dropout param * Minor fixes in training * Fix weights file extension in inference + coverage - Remove unused scripts - More tests - Fixed weights type in inference * Run all hooks * Fix inference testing * Changed anisotropy calculation * Fixed aniso correction and CRF interaction * Remove duplicate tests * Finish rebase + changed step to auto in spinbox * Updated based on feedback from CYHSM Co-Authored-By: Markus Frey <[email protected]> * Added minimal WNet notebook for training * Remove dask * WNet model docs * Added QoL shape info for layer selecter * WNet fixes + PR feedback improvements * Added imagecodecs to open external datasets --------- Co-authored-by: gityves <[email protected]> Co-authored-by: Nevexios <[email protected]> Co-authored-by: Markus Frey <[email protected]>
AdaptiveMotorControlLab · Jul 12, 2023 · 634f5a4 · 634f5a4
1 parent e00806c
commit 634f5a4
Show file tree

Hide file tree

Showing 76 changed files with 4,502 additions and 1,733 deletions.
diff --git a/.coveragerc b/.coveragerc
@@ -0,0 +1,7 @@
+[report]
+exclude_lines =
+    if __name__ == .__main__.:
+
+[run]
+omit =
+    napari_cellseg3d/setup.py
diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
@@ -7,15 +7,11 @@ on:
   push:
     branches:
       - main
-      - npe2
-      - cy/voronoi-otsu
     tags:
       - "v*" # Push events to matching v*, i.e. v1.0, v20.15.10
   pull_request:
     branches:
       - main
-      - npe2
-      - cy/voronoi-otsu
   workflow_dispatch:
 
 jobs:
@@ -55,6 +51,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install setuptools tox tox-gh-actions
+#         pip install git+https://github.com/lucasb-eyer/pydensecrf.git@master#egg=pydensecrf
 
 # this runs the platform-specific tests declared in tox.ini
       - name: Test with tox

diff --git a/.gitignore b/.gitignore
@@ -104,9 +104,11 @@ notebooks/csv_cell_plot.html
 notebooks/full_plot.html
 *.csv
 *.png
+notebooks/instance_test.ipynb
 *.prof
 
 #include test data
 !napari_cellseg3d/_tests/res/test.tif
 !napari_cellseg3d/_tests/res/test.png
 !napari_cellseg3d/_tests/res/test_labels.tif
+cov.syspath.txt
diff --git a/docs/index.rst b/docs/index.rst
@@ -39,8 +39,8 @@ Welcome to napari-cellseg3d's documentation!
     res/code/plugin_convert
     res/code/plugin_metrics
     res/code/model_framework
-    res/code/model_workers
-    res/code/model_instance_seg
+    res/code/workers
+    res/code/instance_segmentation
     res/code/plugin_model_inference
     res/code/plugin_model_training
     res/code/utils

diff --git a/docs/res/code/instance_segmentation.rst b/docs/res/code/instance_segmentation.rst
@@ -0,0 +1,53 @@
+instance_segmentation.py
+===========================================
+
+Classes
+-------------
+
+InstanceMethod
+**************************************
+.. autoclass:: napari_cellseg3d.code_models.instance_segmentation::InstanceMethod
+    :members: __init__
+
+ConnectedComponents
+**************************************
+.. autoclass:: napari_cellseg3d.code_models.instance_segmentation::ConnectedComponents
+    :members: __init__
+
+Watershed
+**************************************
+.. autoclass:: napari_cellseg3d.code_models.instance_segmentation::Watershed
+    :members: __init__
+
+VoronoiOtsu
+**************************************
+.. autoclass:: napari_cellseg3d.code_models.instance_segmentation::VoronoiOtsu
+    :members: __init__
+
+
+Functions
+-------------
+
+binary_connected
+**************************************
+.. autofunction:: napari_cellseg3d.code_models.instance_segmentation::binary_connected
+
+binary_watershed
+**************************************
+.. autofunction:: napari_cellseg3d.code_models.instance_segmentation::binary_watershed
+
+volume_stats
+**************************************
+.. autofunction:: napari_cellseg3d.code_models.instance_segmentation::volume_stats
+
+clear_small_objects
+**************************************
+.. autofunction:: napari_cellseg3d.code_models.instance_segmentation::clear_small_objects
+
+to_instance
+**************************************
+.. autofunction:: napari_cellseg3d.code_models.instance_segmentation::to_instance
+
+to_semantic
+**************************************
+.. autofunction:: napari_cellseg3d.code_models.instance_segmentation::to_semantic
diff --git a/docs/res/code/model_instance_seg.rst b/docs/res/code/model_instance_seg.rst
diff --git a/docs/res/code/plugin_convert.rst b/docs/res/code/plugin_convert.rst
@@ -28,18 +28,3 @@ ThresholdUtils
 **********************************
 .. autoclass:: napari_cellseg3d.code_plugins.plugin_convert::ThresholdUtils
     :members: __init__
-
-Functions
------------------------------------
-
-save_folder
-*****************************************
-.. autofunction:: napari_cellseg3d.code_plugins.plugin_convert::save_folder
-
-save_layer
-****************************************
-.. autofunction:: napari_cellseg3d.code_plugins.plugin_convert::save_layer
-
-show_result
-****************************************
-.. autofunction:: napari_cellseg3d.code_plugins.plugin_convert::show_result
diff --git a/docs/res/code/plugin_model_training.rst b/docs/res/code/plugin_model_training.rst
@@ -18,6 +18,5 @@ Methods
 
 Attributes
 *********************
-
 .. autoclass:: napari_cellseg3d.code_plugins.plugin_model_training::Trainer
    :members:  _viewer, worker, loss_dict, canvas, train_loss_plot, dice_metric_plot
diff --git a/docs/res/code/utils.rst b/docs/res/code/utils.rst
@@ -62,7 +62,3 @@ denormalize_y
 load_images
 **************************************
 .. autofunction:: napari_cellseg3d.utils::load_images
-
-format_Warning
-**************************************
-.. autofunction:: napari_cellseg3d.utils::format_Warning
diff --git a/docs/res/code/model_workers.rst → docs/res/code/workers.rst b/docs/res/code/model_workers.rst → docs/res/code/workers.rst
@@ -1,4 +1,4 @@
-model_workers.py
+workers.py
 ===========================================
 
 
@@ -10,7 +10,7 @@ Class : LogSignal
 
 Attributes
 ************************
-.. autoclass:: napari_cellseg3d.code_models.model_workers::LogSignal
+.. autoclass:: napari_cellseg3d.code_models.workers::LogSignal
    :members: log_signal
    :noindex:
 
@@ -24,7 +24,7 @@ Class : InferenceWorker
 
 Methods
 ************************
-.. autoclass:: napari_cellseg3d.code_models.model_workers::InferenceWorker
+.. autoclass:: napari_cellseg3d.code_models.workers::InferenceWorker
    :members: __init__, log, create_inference_dict, inference
    :noindex:
 
@@ -39,6 +39,6 @@ Class : TrainingWorker
 
 Methods
 ************************
-.. autoclass:: napari_cellseg3d.code_models.model_workers::TrainingWorker
+.. autoclass:: napari_cellseg3d.code_models.workers::TrainingWorker
    :members: __init__, log, train
    :noindex:
diff --git a/docs/res/guides/custom_model_template.rst b/docs/res/guides/custom_model_template.rst
@@ -3,35 +3,33 @@
 Advanced : Declaring a custom model
 =============================================
 
-To add a custom model, you will need a **.py** file with the following structure to be placed in the *napari_cellseg3d/models* folder:
+.. warning::
+    **WIP** : Adding new models is still a work in progress and will likely not work out of the box, leading to errors.
 
-.. note::
-    **WIP** : Currently you must modify :ref:`model_framework.py` as well : import your model class and add it to the ``model_dict`` attribute
-
-::
-
-    def get_net():
-        return ModelClass # should return the class of the model,
-        # for example SegResNet or UNET
+    Please `file an issue`_ if you would like to add a custom model and we will help you get it working.
 
+To add a custom model, you will need a **.py** file with the following structure to be placed in the *napari_cellseg3d/models* folder::
 
-    def get_weights_file():
-        return "weights_file.pth" # name of the weights file for the model,
-        # which should be in *napari_cellseg3d/models/pretrained*
+    class ModelTemplate_(ABC): # replace ABC with your PyTorch model class name
+        use_default_training = True  # not needed for now, will serve for WNet training if added to the plugin
+        weights_file = (
+            "model_template.pth"  # specify the file name of the weights file only
+        ) # download URL goes in pretrained_models.json
 
+        @abstractmethod
+        def __init__(
+            self, input_image_size, in_channels=1, out_channels=1, **kwargs
+        ):
+            """Reimplement this as needed; only include input_image_size if necessary. For now only in/out channels = 1 is supported."""
+            pass
 
-    def get_output(model, input):
-        out = model(input) # should return the model's output as [C, N, D,H,W]
-        # (C: channel, N, batch size, D,H,W : depth, height, width)
-        return out
+        @abstractmethod
+        def forward(self, x):
+            """Reimplement this as needed. Ensure that output is a torch tensor with dims (batch, channels, z, y, x)."""
+            pass
 
 
-    def get_validation(model, val_inputs):
-        val_outputs = model(val_inputs) # should return the proper type for validation
-        # with sliding_window_inference from MONAI
-        return val_outputs
-
+.. note::
+    **WIP** : Currently you must modify :ref:`model_framework.py` as well : import your model class and add it to the ``model_dict`` attribute
 
-    def ModelClass(x1,x2...):
-        # your Pytorch model here...
-        return results # should return as [C, N, D,H,W]
+.. _file an issue: https://github.com/AdaptiveMotorControlLab/CellSeg3d/issues
diff --git a/docs/res/guides/detailed_walkthrough.rst b/docs/res/guides/detailed_walkthrough.rst
@@ -1,6 +1,6 @@
 .. _detailed_walkthrough:
 
-Detailed walkthrough
+Detailed walkthrough - Supervised learning
 ===================================
 
 The following guide will show you how to use the plugin's workflow, starting from human-labeled annotation volume, to running inference on novel volumes.
@@ -109,7 +109,7 @@ of two no matter the size you choose. For optimal performance, make sure to use
 a power of two still, such as 64 or 120.
 
 .. important::
-    Using a too large value for the size will cause memory issues. If this happens, restart napari (better handling for these situations might be added in the future).
+    Using a too large value for the size will cause memory issues. If this happens, restart the worker with smaller volumes.
 
 You also have the option to use data augmentation, which can improve performance and generalization.
 In most cases this should left enabled.

diff --git a/docs/res/guides/inference_module_guide.rst b/docs/res/guides/inference_module_guide.rst
@@ -7,8 +7,9 @@ This module allows you to use pre-trained segmentation algorithms (written in Py
 to automatically label cells.
 
 .. important::
-    Currently, only inference on **3D volumes is supported**. Your image and label folders should both contain a set of
-    **3D image files**, currently either **.tif** or **.tiff**.
+    Currently, only inference on **3D volumes is supported**. If using folders, your images and labels folders
+    should both contain a set of **3D image files**, either **.tif** or **.tiff**.
+    Otherwise you may run inference on layers in napari.
 
 Currently, the following pre-trained models are available :
 
@@ -20,13 +21,18 @@ SegResNet        `3D MRI brain tumor segmentation using autoencoder regularizati
 TRAILMAP_MS       A PyTorch implementation of the `TRAILMAP project on GitHub`_ pretrained with mesoSPIM data
 TRAILMAP          An implementation of the `TRAILMAP project on GitHub`_ using a `3DUNet for PyTorch`_
 SwinUNetR         `Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images`_
+WNet             `WNet, A Deep Model for Fully Unsupervised Image Segmentation`_
 ==============   ================================================================================================
 
 .. _Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation: https://arxiv.org/pdf/1606.04797.pdf
 .. _3D MRI brain tumor segmentation using autoencoder regularization: https://arxiv.org/pdf/1810.11654.pdf
 .. _TRAILMAP project on GitHub: https://github.com/AlbertPun/TRAILMAP
 .. _3DUnet for Pytorch: https://github.com/wolny/pytorch-3dunet
 .. _Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images: https://arxiv.org/abs/2201.01266
+.. _WNet, A Deep Model for Fully Unsupervised Image Segmentation: https://arxiv.org/abs/1711.08506
+
+.. note::
+    For WNet-specific instruction please refer to  the appropriate section below.
 
 Interface and functionalities
 --------------------------------
@@ -67,8 +73,7 @@ Interface and functionalities
 
 * **Instance segmentation** :
 
-  | You can convert the semantic segmentation into instance labels by using either the `watershed`_ or `connected components`_ method.
-  | You can set the probability threshold from which a pixel is considered as a valid instance, as well as the minimum size in pixels for objects. All smaller objects will be removed.
+  | You can convert the semantic segmentation into instance labels by using either the Voronoi-Otsu, `Watershed`_ or `Connected Components`_ method, as detailed in :ref:`utils_module_guide`.
   | Instance labels will be saved (and shown if applicable) separately from other results.
 
 
@@ -78,7 +83,7 @@ Interface and functionalities
 
 * **Computing objects statistics** :
 
-    You can choose to compute various stats from the labels and save them to a csv for later use.
+    You can choose to compute various stats from the labels and save them to a .csv for later use.
 
     This includes, for each object :
 
@@ -98,13 +103,6 @@ Interface and functionalities
 
     In the ``notebooks`` folder you can find an example of plotting cell statistics using the result csv.
 
-* **Viewing results** :
-
-  | You can also select whether you'd like to **see the results** in napari afterwards.
-  | By default the first image processed will be displayed, but you can choose to display up to **ten at once**.
-  | You can also request to see the originals.
-
-
 When you are done choosing your parameters, you can press the **Start** button to begin the inference process.
 Once it has finished, results will be saved then displayed in napari; each output will be paired with its original.
 On the left side, a progress bar and a log will keep you informed on the process.
@@ -115,7 +113,7 @@ On the left side, a progress bar and a log will keep you informed on the process
     |    ``{original_name}_{model}_{date & time}_pred{id}.file_ext``
     | For example, using a VNet on the third image of a folder, called "somatomotor.tif" will yield the following name :
     |   *somatomotor_VNet_2022_04_06_15_49_42_pred3.tif*
-    | Instance labels will have the "Instance_seg" prefix appened to the name.
+    | Instance labels will have the "Instance_seg" prefix appended to the name.
 
 
 .. hint::
@@ -128,8 +126,19 @@ On the left side, a progress bar and a log will keep you informed on the process
 .. note::
     You can save the log after the worker is finished to easily remember which parameters you ran inference with.
 
+WNet
+--------------------------------
+
+The WNet model, from the paper `WNet, A Deep Model for Fully Unsupervised Image Segmentation`_, is a fully unsupervised model that can be used to segment images without any labels.
+It clusters pixels based on brightness, and can be used to segment cells in a variety of modalities.
+Its use and available options are similar to the above models, with a few differences :
+.. note::
+    | Our provided, pre-trained model should use an input size of 64x64x64. As such, window inference is always enabled
+    | and set to 64. If you want to use a different size, you will have to train your own model using the provided notebook.
+All it requires are images; for nucleus segmentation, it is recommended to use 2 classes (default).
+
 Source code
 --------------------------------
 * :doc:`../code/plugin_model_inference`
 * :doc:`../code/model_framework`
-* :doc:`../code/model_workers`
+* :doc:`../code/workers`