Started docs update

Author: C-Achard

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/test_training.py

@@ -124,7 +124,7 @@ def test_unsupervised_training(make_napari_viewer_proxy):
         {"image": im_path_str, "label": im_path_str}
     ]
     widget.worker._get_data()
-    eval_res = widget.worker._eval(
+    eval_res = widget.worker.eval(
         model=WNetFixture(),
         epoch=-10,
     )

napari_cellseg3d/code_models/worker_training.py

@@ -108,10 +108,10 @@ def set_download_log(self, widget):
         self.downloader.log_widget = widget
 
     def log(self, text):
-        """Sends a signal that ``text`` should be logged
+        """Sends a Qt signal that the provided text should be logged
         Goes in a Log object, defined in :py:mod:`napari_cellseg3d.interface
         Sends a signal to the main thread to log the text.
-        Signal is defined in napari_cellseg3d.workers_utils.LogSignal
+        Signal is defined in napari_cellseg3d.workers_utils.LogSignal.
 
         Args:
             text (str): text to logged
@@ -653,7 +653,7 @@ def train(
                 ):
                     model.eval()
                     self.log("Validating...")
-                    yield self._eval(model, epoch)  # validation
+                    yield self.eval(model, epoch)  # validation
 
                     if self._abort_requested:
                         self.dataloader = None
@@ -736,7 +736,7 @@ def train(
             self.quit()
             raise e
 
-    def _eval(self, model, epoch) -> TrainingReport:
+    def eval(self, model, epoch) -> TrainingReport:
         with torch.no_grad():
             device = self.config.device
             for _k, val_data in enumerate(self.eval_dataloader):