added learning rate param + number of cells

Author: C-Achard

src/napari_cellseg3d/model_framework.py

@@ -3,6 +3,7 @@
 
 import napari
 import torch
+
 # Qt
 from qtpy.QtWidgets import QLineEdit
 from qtpy.QtWidgets import QProgressBar

src/napari_cellseg3d/model_workers.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import torch
+
 # MONAI
 from monai.data import CacheDataset
 from monai.data import DataLoader
@@ -30,6 +31,7 @@
 from monai.transforms import Zoom
 from napari.qt.threading import GeneratorWorker
 from napari.qt.threading import WorkerBaseSignals
+
 # Qt
 from qtpy.QtCore import Signal
 from tifffile import imwrite
@@ -378,6 +380,7 @@ def __init__(
         data_dicts,
         max_epochs,
         loss_function,
+        learning_rate,
         val_interval,
         batch_size,
         results_path,
@@ -401,6 +404,8 @@ def __init__(
 
             * loss_function : the loss function to use for training
 
+            * learning_rate : the learning rate of the optimizer
+
             * val_interval : the interval at which to perform validation (e.g. if 2 will validate once every 2 epochs.) Also determines frequency of saving, depending on whether the metric is better or not
 
             * batch_size : the batch size to use for training
@@ -428,6 +433,7 @@ def __init__(
         self.data_dicts = data_dicts
         self.max_epochs = max_epochs
         self.loss_function = loss_function
+        self.learning_rate = learning_rate
         self.val_interval = val_interval
         self.batch_size = batch_size
         self.results_path = results_path
@@ -486,6 +492,8 @@ def train(self):
 
         * loss_function : the loss function to use for training
 
+        * learning rate : the learning rate of the optimizer
+
         * val_interval : the interval at which to perform validation (e.g. if 2 will validate once every 2 epochs.) Also determines frequency of saving, depending on whether the metric is better or not
 
         * batch_size : the batch size to use for training
@@ -653,7 +661,7 @@ def train(self):
         )
         print("\nDone")
 
-        optimizer = torch.optim.Adam(model.parameters(), 1e-3)
+        optimizer = torch.optim.Adam(model.parameters(), self.learning_rate)
         dice_metric = DiceMetric(include_background=True, reduction="mean")
 
         best_metric = -1

src/napari_cellseg3d/plugin_dock.py

@@ -1,9 +1,11 @@
 import os
 import warnings
+
 # import shutil
 from pathlib import Path
 
 import pandas as pd
+
 # Qt
 from qtpy.QtWidgets import QVBoxLayout
 from qtpy.QtWidgets import QWidget

src/napari_cellseg3d/plugin_helper.py

@@ -1,4 +1,5 @@
 import napari
+
 # Qt
 from qtpy.QtWidgets import QSizePolicy
 from qtpy.QtWidgets import QVBoxLayout

src/napari_cellseg3d/plugin_metrics.py

@@ -157,8 +157,8 @@ def remove_plots(self):
 
     def compute_dice(self):
         """Computes the dice metric between pairs of labels. Rotates the prediction label to find matching orientation as well."""
-        u = 0
-        t = 0
+        # u = 0
+        # t = 0
         total_metrics = []
         self.canvas = (
             None  # kind of unsafe way to stack plots... but it works.
@@ -205,12 +205,29 @@ def compute_dice(self):
                     f"Padded sizes of images do not match ! Padded ground label : {ground.shape} Padded pred label : {pred.shape}"
                 )
             # if u < 1:
-            self._viewer.add_image(ground, name="ground", colormap="blue",opacity=0.7)
-            self._viewer.add_image(pred, name="pred", colormap="red")
-            self._viewer.add_image(np.rot90(pred[0][0], axes=(0,1)), name="pred flip 0", colormap="red",opacity=0.7)
-            self._viewer.add_image(np.rot90(pred[0][0], axes=(1,2)), name="pred flip 1", colormap="red",opacity=0.7)
-            self._viewer.add_image(np.rot90(pred[0][0], axes=(0,2)), name="pred flip 2", colormap="red",opacity=0.7)
-            u+=1
+            # self._viewer.add_image(
+            #     ground, name="ground", colormap="blue", opacity=0.7
+            # )
+            # self._viewer.add_image(pred, name="pred", colormap="red")
+            # self._viewer.add_image(
+            #     np.rot90(pred[0][0], axes=(0, 1)),
+            #     name="pred flip 0",
+            #     colormap="red",
+            #     opacity=0.7,
+            # )
+            # self._viewer.add_image(
+            #     np.rot90(pred[0][0], axes=(1, 2)),
+            #     name="pred flip 1",
+            #     colormap="red",
+            #     opacity=0.7,
+            # )
+            # self._viewer.add_image(
+            #     np.rot90(pred[0][0], axes=(0, 2)),
+            #     name="pred flip 2",
+            #     colormap="red",
+            #     opacity=0.7,
+            # )
+            # u += 1
 
             # TODO add rotation toggle
             pred_flip_x = np.rot90(pred[0][0], axes=(0, 1))
@@ -224,10 +241,15 @@ def compute_dice(self):
                 for i in range(3):
                     scores.append(utils.dice_coeff(np.flip(p, axis=i), ground))
 
-            if t <1 :
-                for i in range(3):
-                    self._viewer.add_image(np.flip(pred_flip_x,axis=i), name=f"flip", colormap="green",opacity=0.7)
-                t+=1
+            # if t < 1:
+            #     for i in range(3):
+            #         self._viewer.add_image(
+            #             np.flip(pred_flip_x, axis=i),
+            #             name=f"flip",
+            #             colormap="green",
+            #             opacity=0.7,
+            #         )
+            #     t += 1
 
             # print(scores)
             score = max(scores)

src/napari_cellseg3d/plugin_model_inference.py

@@ -3,6 +3,7 @@
 import numpy as np
 
 import napari
+
 # Qt
 from qtpy.QtWidgets import QLabel
 from qtpy.QtWidgets import QSizePolicy
@@ -626,7 +627,9 @@ def on_yield(data, widget):
 
             if data["instance_labels"] is not None:
 
-                widget.log.print_and_log(f"\nNUMBER OF CELLS : {np.amax(data['instance_labels'])}\n")
+                widget.log.print_and_log(
+                    f"\nNUMBER OF CELLS : {np.amax(data['instance_labels'])}\n"
+                )
 
                 name = f"instance_labels_{image_id}"
                 instance_layer = viewer.add_labels(

src/napari_cellseg3d/plugin_model_training.py

@@ -10,13 +10,15 @@
     FigureCanvasQTAgg as FigureCanvas,
 )
 from matplotlib.figure import Figure
+
 # MONAI
 from monai.losses import DiceCELoss
 from monai.losses import DiceFocalLoss
 from monai.losses import DiceLoss
 from monai.losses import FocalLoss
 from monai.losses import GeneralizedDiceLoss
 from monai.losses import TverskyLoss
+
 # Qt
 from qtpy.QtWidgets import QLabel
 from qtpy.QtWidgets import QProgressBar
@@ -163,6 +165,7 @@ def __init__(
         """At which epochs to perform validation. E.g. if 2, will run validation on epochs 2,4,6..."""
         self.patch_size = []
         """The size of samples to be extracted from images"""
+        self.learning_rate = 1e-3
 
         self.model = None  # TODO : custom model loading ?
         self.worker = None
@@ -222,6 +225,20 @@ def __init__(
         )
         self.lbl_val_interv_choice = QLabel("Validation interval : ", self)
 
+        self.learning_rate_dict = {
+            "1e-3": 1e-3,
+            "1e-4": 1e-4,
+            "1e-5": 1e-5,
+            "1e-6": 1e-6,
+        }
+
+        (
+            self.learning_rate_choice,
+            self.lbl_learning_rate_choice,
+        ) = ui.make_combobox(
+            self.learning_rate_dict.keys(), label="Learning rate"
+        )
+
         self.augment_choice = ui.make_checkbox("Augment data")
 
         # TODO add self.tabs, self.close_buttons etc...
@@ -528,8 +545,20 @@ def build(self):
                 r=5,
                 b=5,
             ),
-            alignment=ui.LEFT_AL,
+            # alignment=ui.LEFT_AL,
         )  # batch size
+        train_param_group_l.addWidget(
+            ui.combine_blocks(
+                self.learning_rate_choice,
+                self.lbl_learning_rate_choice,
+                min_spacing=spacing,
+                horizontal=False,
+                l=5,
+                t=5,
+                r=5,
+                b=5,
+            )
+        )  # learning rate
         train_param_group_l.addWidget(
             ui.combine_blocks(
                 self.epoch_choice,
@@ -676,6 +705,8 @@ def start(self):
             self.data = self.create_train_dataset_dict()
             self.max_epochs = self.epoch_choice.value()
 
+            self.learning_rate = self.learning_rate_dict[self.learning_rate_choice.currentText()]
+
             self.patch_size = []
             [
                 self.patch_size.append(w.value())
@@ -715,6 +746,7 @@ def start(self):
                 data_dicts=self.data,
                 max_epochs=self.max_epochs,
                 loss_function=self.get_loss(self.loss_choice.currentText()),
+                learning_rate=self.learning_rate,
                 val_interval=self.val_interval,
                 batch_size=self.batch_size,
                 results_path=self.results_path,
@@ -804,7 +836,13 @@ def on_yield(data, widget):
         widget.update_loss_plot(data["losses"], data["val_metrics"])
 
         if widget.stop_requested:
-            torch.save(data["weights"], os.path.join(widget.results_path, f"latest_weights_aborted_training_{utils.get_date_time()}.pth"))
+            torch.save(
+                data["weights"],
+                os.path.join(
+                    widget.results_path,
+                    f"latest_weights_aborted_training_{utils.get_date_time()}.pth",
+                ),
+            )
             widget.stop_requested = False
 
     # def clean_cache(self):

src/napari_cellseg3d/plugin_review.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pims
 import skimage.io as io
+
 # Qt
 from qtpy import QtGui
 from qtpy.QtWidgets import QLabel

src/napari_cellseg3d/plugin_utilities.py

@@ -1,8 +1,10 @@
 import napari
+
 # Qt
 from qtpy.QtWidgets import QTabWidget
 
 from napari_cellseg3d.plugin_convert import ConvertUtils
+
 # local
 from napari_cellseg3d.plugin_crop import Cropping
 from napari_cellseg3d.plugin_metrics import MetricsUtils