WNet training (#47)

* Add LayerNorm

* Change softmax arg

* Num group 2

* Update model.py

* Update model.py

* Reduce depth of WNet

* Started WNet training UI

* Workable WNet training prototype

* Fixes

* Test fixes

* Temp fix for CRF (#46)

* Minor fixes

* Tests & training

* Fix tests + new weights

* Fix ETA precision

* Docstring update

* Update plugin_model_training.py

* Update contrast limit when updating layers

* Update config.py

* Fixed normalization

* Update plugin_model_training.py

* Update workers_utils.py

* Trying to fix input normalization

* Fix name mismatch

* Fix decoder evaluation

* Update dice calculation

* Update dice coeff

* Update worker_training.py

* Fix eval detach

* Fix Dice list for WNet

* Updated validation UI

* Tooltips and show_results update

* Plots update

* Plot + log_parameters

* Update worker_training.py

* Disable WANDB for now + log param tweaks

* UI/log tweaks

* Functional WNet training

* Clean exit / free memory attempt

* Cleanup + tests

- Removed previous train script
- Fix tests
- Enable test workflow on GH

* Deploy memory usage fix in inference as well

* Memory usage fix

* UI tweak

* WNet cleanup + supervised training improvements

* Change Dice metric include_background for WNet

To avoid Max Dice calculation

* Set better default LR across un/supervised

* Update model.py

* Update WNet weights

* Fix default LR + sup. test

* Fix new unsup LR in tests

* Fix dir for saving in tests

* Testing fixes

Due to Singleton Trainer widget

* Test unsupervised training and raise coverage

* WNet eval test

* Fix order for model deletion

* Extend supervised train tests

* Started docs update

* Update plugin_model_training.py

* Fixed filepaths

* Fix paths in test (use pathlib)

* Updated workers config

* Fixed parse_default_path test

* Ignore wandb results in gitignore

* Enable GH Actions tests on branch temporarily

* Fixed deletion of Qt imports in interface

* Reverted include_background=True in Dice

* Reintroduced best Dice channel seeking + refacto

* Improve filepath messages

* Fix unsup image loading when not validating

* Fix training tests

Author: C-Achard

.github/workflows/test_and_deploy.yml

@@ -7,7 +7,7 @@ on:
   push:
     branches:
       - main
-      - cy/utils
+      - cy/wnet-train
     tags:
       - "v*" # Push events to matching v*, i.e. v1.0, v20.15.10
   pull_request:

.gitignore

@@ -103,6 +103,7 @@ venv/
 /docs/res/logo/old_logo/
 /reqs/
 /loss_plots/
+/wandb/
 notebooks/csv_cell_plot.html
 notebooks/full_plot.html
 *.csv

docs/res/code/model_framework.rst

@@ -12,7 +12,7 @@ Class : ModelFramework
 Methods
 **********************
 .. autoclass:: napari_cellseg3d.code_models.model_framework::ModelFramework
-   :members:  __init__, send_log, save_log, save_log_to_path, display_status_report, create_train_dataset_dict, get_model, get_available_models, get_device, empty_cuda_cache
+   :members:  __init__, send_log, save_log, save_log_to_path, display_status_report, create_train_dataset_dict, get_available_models, get_device, empty_cuda_cache
    :noindex:
 
 

docs/res/code/plugin_model_training.rst

@@ -11,12 +11,12 @@ Class : Trainer
 Methods
 **********************
 .. autoclass:: napari_cellseg3d.code_plugins.plugin_model_training::Trainer
-   :members:  __init__, get_loss, check_ready, send_log, start, on_start, on_finish, on_error, on_yield, plot_loss, update_loss_plot
+   :members:  __init__, check_ready, send_log, start, on_start, on_finish, on_error, on_yield, update_loss_plot
    :noindex:
 
 
 
 Attributes
 *********************
 .. autoclass:: napari_cellseg3d.code_plugins.plugin_model_training::Trainer
-   :members:  _viewer, worker, loss_dict, canvas, train_loss_plot, dice_metric_plot
+   :members:  _viewer, worker, canvas

docs/res/code/workers.rst

@@ -10,7 +10,7 @@ Class : LogSignal
 
 Attributes
 ************************
-.. autoclass:: napari_cellseg3d.code_models.workers::LogSignal
+.. autoclass:: napari_cellseg3d.code_models.workers_utils::LogSignal
    :members: log_signal
    :noindex:
 
@@ -24,21 +24,47 @@ Class : InferenceWorker
 
 Methods
 ************************
-.. autoclass:: napari_cellseg3d.code_models.workers::InferenceWorker
+.. autoclass:: napari_cellseg3d.code_models.worker_inference::InferenceWorker
    :members: __init__, log, create_inference_dict, inference
    :noindex:
 
 .. _here: https://napari-staging-site.github.io/guides/stable/threading.html
 
 
-Class : TrainingWorker
+Class : TrainingWorkerBase
 -------------------------------------------
 
 .. important::
     Inherits from :py:class:`napari.qt.threading.GeneratorWorker`
 
 Methods
 ************************
-.. autoclass:: napari_cellseg3d.code_models.workers::TrainingWorker
+.. autoclass:: napari_cellseg3d.code_models.worker_training::TrainingWorkerBase
    :members: __init__, log, train
    :noindex:
+
+
+Class : WNetTrainingWorker
+-------------------------------------------
+
+.. important::
+    Inherits from :py:class:`TrainingWorkerBase`
+
+Methods
+************************
+.. autoclass:: napari_cellseg3d.code_models.worker_training::WNetTrainingWorker
+   :members: __init__, train, eval, get_patch_dataset, get_dataset_eval, get_dataset
+   :noindex:
+
+
+Class : SupervisedTrainingWorker
+-------------------------------------------
+
+.. important::
+    Inherits from :py:class:`TrainingWorkerBase`
+
+Methods
+************************
+.. autoclass:: napari_cellseg3d.code_models.worker_training::SupervisedTrainingWorker
+   :members: __init__, train
+   :noindex:

docs/res/guides/detailed_walkthrough.rst

@@ -120,9 +120,9 @@ Finally, the last tab lets you choose :
 
     * SegResNet is a lightweight model (low memory requirements) from MONAI originally designed for 3D fMRI data.
     * VNet is a larger (than SegResNet) CNN from MONAI designed for medical image segmentation.
-    * TRAILMAP is our PyTorch implementation of a 3D CNN model trained for axonal detection in cleared tissue.
     * TRAILMAP_MS is our implementation in PyTorch additionally trained on mouse cortical neural nuclei from mesoSPIM data.
-    * Note, the code is very modular, so it is relatively straightforward to use (and contribute) your model as well.
+    * SwinUNetR is a MONAI implementation of the SwinUNetR model. It is costly in compute and memory, but can achieve high performance.
+    * WNet is our reimplementation of an unsupervised model, which can be used to produce segmentation without labels.
 
 
 * The loss : for object detection in 3D volumes you'll likely want to use the Dice or Dice-focal Loss.
@@ -239,13 +239,12 @@ Scoring, review, analysis
 ----------------------------
 
 
-.. Using the metrics utility module, you can compare the model's predictions to any ground truth
-labels you might have.
-Simply provide your prediction and ground truth labels, and compute the results.
-A Dice metric of 1 indicates perfect matching, whereas a score of 0 indicates complete mismatch.
-Select which score **you consider as sub-optimal**, and all results below this will be **shown in napari**.
-If at any time the **orientation of your prediction labels changed compared to the ground truth**, check the
-"Find best orientation" option to compensate for it.
+.. Using the metrics utility module, you can compare the model's predictions to any ground truth labels you might have.
+    Simply provide your prediction and ground truth labels, and compute the results.
+    A Dice metric of 1 indicates perfect matching, whereas a score of 0 indicates complete mismatch.
+    Select which score **you consider as sub-optimal**, and all results below this will be **shown in napari**.
+    If at any time the **orientation of your prediction labels changed compared to the ground truth**, check the
+    "Find best orientation" option to compensate for it.
 
 
 Labels review

docs/res/guides/training_module_guide.rst

@@ -4,7 +4,7 @@ Training module guide - Unsupervised models
 ==============================================
 
 .. important::
-    The WNet training is for now only available in the provided jupyter notebook, in the ``notebooks`` folder.
+    The WNet training is for now available as part of the plugin in the Training module.
     Please see the :ref:`training_wnet` section for more information.
 
 Training module guide - Supervised models
@@ -25,14 +25,15 @@ Model            Link to original paper
 ==============   ================================================================================================
 VNet             `Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation`_
 SegResNet        `3D MRI brain tumor segmentation using autoencoder regularization`_
-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`_
+TRAILMAP_MS       An implementation of the `TRAILMAP project on GitHub`_ using `3DUNet for PyTorch`_
+SwinUNetR         `Swin UNETR, Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images`_
 ==============   ================================================================================================
 
 .. _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 UNETR, Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images: https://arxiv.org/abs/2201.01266
 
 .. important::
     | The machine learning models used by this program require all images of a dataset to be of the same size.

docs/res/guides/training_wnet.rst

@@ -15,21 +15,45 @@ the model was trained on; you can retrain from our pretrained model to your set
 The model has two losses, the SoftNCut loss which clusters pixels according to brightness, and a reconstruction loss, either
 Mean Square Error (MSE) or Binary Cross Entropy (BCE).
 Unlike the original paper, these losses are added in a weighted sum and the backward pass is performed for the whole model at once.
-The SoftNcuts is bounded between 0 and 1; the MSE may take large values.
+The SoftNcuts is bounded between 0 and 1; the MSE may take large positive values.
 
-For good performance, one should wait for the SoftNCut to reach a plateau, the reconstruction loss must also diminish but it's generally less critical.
+For good performance, one should wait for the SoftNCut to reach a plateau; the reconstruction loss must also diminish but it's generally less critical.
 
+Parameters
+-------------------------------
+
+When using the WNet training module, additional options will be provided in the Advanced tab of the training module:
+
+- Number of classes : number of classes to segment (default 2). Additional classes will result in a more progressive segmentation according to brightness; can be useful if you have "halos" around your objects or artifacts with a significantly different brightness.
+- Reconstruction loss : either MSE or BCE (default MSE). MSE is more sensitive to outliers, but can be more precise; BCE is more robust to outliers but can be less precise.
+
+- NCuts parameters:
+    - Intensity sigma : standard deviation of the feature similarity term (brightness here, default 1)
+    - Spatial sigma : standard deviation of the spatial proximity term (default 4)
+    - Radius : radius of the loss computation in pixels (default 2)
+
+.. note::
+    Intensity sigma depends on pixel values in the image. The default of 1 is tailored to images being mapped between 0 and 100, which is done automatically by the plugin.
+.. note::
+    Raising the radius might improve performance in some cases, but will also greatly increase computation time.
+
+- Weights for the sum of losses :
+    - NCuts weight : weight of the NCuts loss (default 0.5)
+    - Reconstruction weight : weight of the reconstruction loss (default 0.5*1e-2)
+
+.. note::
+    The weight of the reconstruction loss should be adjusted according to its empirical value; ideally the reconstruction loss should be of the same order of magnitude as the NCuts loss after being multiplied by its weight.
 
 Common issues troubleshooting
 ------------------------------
-If you do not find a satisfactory answer here, please `open an issue`_ !
+If you do not find a satisfactory answer here, please do not hesitate to `open an issue`_ on GitHub.
 
-- **The NCuts loss explodes after a few epochs** : Lower the learning rate
+- **The NCuts loss explodes after a few epochs** : Lower the learning rate, first by a factor of two, then ten.
 
 - **The NCuts loss does not converge and is unstable** :
-  The normalization step might not be adapted to your images. Disable normalization and change intensity_sigma according to the distribution of values in your image; for reference, by default images are remapped to values between 0 and 100, and intensity_sigma=1.
+  The normalization step might not be adapted to your images. Disable normalization and change intensity_sigma according to the distribution of values in your image. For reference, by default images are remapped to values between 0 and 100, and intensity_sigma=1.
 
-- **Reconstruction (decoder) performance is poor** : switch to BCE and set the scaling factor of the reconstruction loss ot 0.5, OR adjust the weight of the MSE loss to make it closer to 1.
+- **Reconstruction (decoder) performance is poor** : switch to BCE and set the scaling factor of the reconstruction loss ot 0.5, OR adjust the weight of the MSE loss to make it closer to 1 in the weighted sum.
 
 
 .. _WNet, A Deep Model for Fully Unsupervised Image Segmentation: https://arxiv.org/abs/1711.08506

napari_cellseg3d/_tests/fixtures.py

@@ -1,3 +1,4 @@
+import torch
 from qtpy.QtWidgets import QTextEdit
 
 from napari_cellseg3d.utils import LOGGER as logger
@@ -17,3 +18,60 @@ def warn(self, warning):
 
     def error(self, e):
         raise (e)
+
+
+class WNetFixture(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mock_conv = torch.nn.Conv3d(1, 1, 1)
+        self.mock_conv.requires_grad_(False)
+
+    def forward_encoder(self, x):
+        return x
+
+    def forward_decoder(self, x):
+        return x
+
+    def forward(self, x):
+        return self.forward_encoder(x), self.forward_decoder(x)
+
+
+class ModelFixture(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mock_conv = torch.nn.Conv3d(1, 1, 1)
+        self.mock_conv.requires_grad_(False)
+
+    def forward(self, x):
+        return x
+
+
+class OptimizerFixture:
+    def __init__(self):
+        self.param_groups = []
+        self.param_groups.append({"lr": 0})
+
+    def zero_grad(self):
+        pass
+
+    def step(self, *args):
+        pass
+
+
+class SchedulerFixture:
+    def step(self, *args):
+        pass
+
+
+class LossFixture:
+    def __call__(self, *args):
+        return self
+
+    def backward(self, *args):
+        pass
+
+    def item(self):
+        return 0
+
+    def detach(self):
+        return self

napari_cellseg3d/_tests/test_inference.py

@@ -23,7 +23,7 @@ def test_onnx_inference(make_napari_viewer_proxy):
     path = str(Path(PRETRAINED_WEIGHTS_DIR).resolve() / "wnet.onnx")
     assert Path(path).is_file()
     dims = 64
-    batch = 2
+    batch = 1
     x = torch.randn(size=(batch, 1, dims, dims, dims))
     worker = ONNXModelWrapper(file_location=path)
     assert worker.eval() is None
@@ -66,19 +66,23 @@ def test_inference_on_folder():
     config.images_filepaths = [
         str(Path(__file__).resolve().parent / "res/test.tif")
     ]
-    config.sliding_window_config.window_size = 64
+
+    config.sliding_window_config.window_size = 8
 
     class mock_work:
-        def __call__(self, x):
-            return x
+        @staticmethod
+        def eval():
+            return True
 
-        def eval(self):
-            return None
+        def __call__(self, x):
+            return torch.Tensor(x)
 
     worker = InferenceWorker(worker_config=config)
     worker.aniso_transform = mock_work()
 
-    image = torch.Tensor(rand_gen.random((1, 1, 64, 64, 64)))
+    image = torch.Tensor(rand_gen.random(size=(1, 1, 8, 8, 8)))
+    assert image.shape == (1, 1, 8, 8, 8)
+    assert image.dtype == torch.float32
     res = worker.inference_on_folder(
         {"image": image},
         0,