various fixes

Fixed saving in training, changed format of output a bit, added warning in pad for values slightly above power of two + updated test

Author: C-Achard

src/napari_cellseg3d/_tests/test_dock_widget.py

@@ -7,7 +7,6 @@
 
 
 def test_prepare(make_napari_viewer):
-    path_to_csv = Path(os.path.dirname(os.path.realpath(__file__)) + "/res")
     path_image = Path(
         os.path.dirname(os.path.realpath(__file__)) + "/res/test.tif"
     )
@@ -16,7 +15,7 @@ def test_prepare(make_napari_viewer):
     viewer.add_image(image)
     widget = Datamanager(viewer)
 
-    widget.prepare(path_to_csv, ".tif", "", False, False)
+    widget.prepare(path_image, ".tif", "", False, False)
 
     assert widget.filetype == ".tif"
     assert widget.as_folder == False

src/napari_cellseg3d/_tests/test_utils.py

@@ -67,6 +67,13 @@ def test_get_padding_dim():
 
     assert pad == [2048, 32, 64]
 
+    tensor = torch.randn(65,70,80)
+    size = tensor.size()
+
+    pad = utils.get_padding_dim(size)
+
+    assert pad == [128,128,128]
+
 
 def test_normalize_x():
 

src/napari_cellseg3d/launch_review.py

@@ -187,7 +187,7 @@ def quicksave():
         #     quicksave()
         #     viewer.window.close()
 
-        return dirname, quicksave() #, quicksave_quit()
+        return dirname, quicksave()  # , quicksave_quit()
 
     # gui = file_widget.show(run=True)  # dirpicker.show(run=True)
 

src/napari_cellseg3d/log_utility.py

@@ -20,8 +20,8 @@ def __init__(self, parent):
 
         self.lock = threading.Lock()
 
-    def receive_log(self, text):
-        self.print_and_log(text)
+    # def receive_log(self, text):
+    #     self.print_and_log(text)
 
     def print_and_log(self, text):
         """Utility used to both print to terminal and log text to a QTextEdit
@@ -31,10 +31,14 @@ def print_and_log(self, text):
             text (str): Text to be printed and logged
 
         """
-        with self.lock:
+        self.lock.acquire()
+        try:
             print(text)
+            # causes issues if you clik on terminal (tied to CMD QuickEdit mode)
             self.moveCursor(QTextCursor.End)
             self.insertPlainText(f"\n{text}")
             self.verticalScrollBar().setValue(
                 self.verticalScrollBar().maximum()
             )
+        finally:
+            self.lock.release()

src/napari_cellseg3d/model_framework.py

@@ -133,13 +133,22 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
     def send_log(self, text):
         self.log.print_and_log(text)
 
-    def save_log(self):
-        """Saves the worker's log to disk at self.results_path when called"""
+    def save_log(self, spec_path=None):
+        """Saves the worker's log to disk at self.results_path when called
+
+        Args:
+            spec_path: if specified, saves to path instead of self.results_path
+        """
         log = self.log.toPlainText()
 
+        if spec_path is None:
+            path = self.results_path
+        else:
+            path = spec_path
+
         if len(log) != 0:
             with open(
-                self.results_path + f"/Log_report_{utils.get_date_time()}.txt",
+                path + f"/Log_report_{utils.get_date_time()}.txt",
                 "x",
             ) as f:
                 f.write(log)

src/napari_cellseg3d/model_workers.py

@@ -200,13 +200,13 @@ def inference(self):
         self.log("\nChecking dimensions...")
         pad = utils.get_padding_dim(check)
         # print(pad)
-        dims =128
-        dims=64
+        dims = 128
+        dims = 64
 
         model = self.model_dict["class"].get_net()
         if self.model_dict["name"] == "SegResNet":
             model = self.model_dict["class"].get_net()(
-                input_image_size=[dims,dims,dims],  # TODO FIX !
+                input_image_size=[dims, dims, dims],  # TODO FIX !
                 out_channels=1,
                 # dropout_prob=0.3,
             )
@@ -511,7 +511,7 @@ def train(self):
 
         * data_dicts : dict from :py:func:`Trainer.create_train_dataset_dict`
 
-        * max_epochs : the amout of epochs to train for
+        * max_epochs : the amount of epochs to train for
 
         * loss_function : the loss function to use for training
 
@@ -757,7 +757,7 @@ def train(self):
                     f"* {step}/{len(train_ds) // train_loader.batch_size}, "
                     f"Train loss: {loss.detach().item():.4f}"
                 )
-                yield {"plot":False, "weights": model.state_dict()}
+                yield {"plot": False, "weights": model.state_dict()}
 
             epoch_loss /= step
             epoch_loss_values.append(epoch_loss)
@@ -820,11 +820,12 @@ def train(self):
                     if metric > best_metric:
                         best_metric = metric
                         best_metric_epoch = epoch + 1
+                        self.log("Saving best metric model")
                         torch.save(
                             model.state_dict(),
                             os.path.join(self.results_path, weights_filename),
                         )
-                        self.log("Saved best metric model")
+                        self.log("Saving complete")
                     self.log(
                         f"Current epoch: {epoch + 1}, Current mean dice: {metric:.4f}"
                         f"\nBest mean dice: {best_metric:.4f} "

src/napari_cellseg3d/plugin_convert.py

@@ -164,7 +164,7 @@ def build(self):
         ui.add_blank(layout=layout, widget=self)
         layout.addWidget(self.lbl_error)
 
-        ui.make_scrollable(layout, self, min_wh=[230, 300], base_wh=[230, 350])
+        ui.make_scrollable(layout, self, min_wh=[230, 400], base_wh=[230, 450])
 
     def folder_to_semantic(self):
         """Converts folder of labels to semantic labels"""

src/napari_cellseg3d/plugin_crop.py

@@ -283,14 +283,14 @@ def add_crop_sliders(
         print(f"Crop variables")
         print(image_stack.shape)
 
-
-
         # define crop sizes and boundaries for the image
         crop_sizes = [self._crop_size_x, self._crop_size_y, self._crop_size_z]
         for i in range(len(crop_sizes)):
             if crop_sizes[i] > image_stack.shape[i]:
                 crop_sizes[i] = image_stack.shape[i]
-                warnings.warn(f"WARNING : Crop dimension in axis {i} was too large at {crop_sizes[i]}, it was set to {image_stack.shape[i]}")
+                warnings.warn(
+                    f"WARNING : Crop dimension in axis {i} was too large at {crop_sizes[i]}, it was set to {image_stack.shape[i]}"
+                )
         cropx, cropy, cropz = crop_sizes
         # shapez, shapey, shapex = image_stack.shape
         ends = np.asarray(image_stack.shape) - np.asarray(crop_sizes) + 1

src/napari_cellseg3d/plugin_dock.py

@@ -162,7 +162,7 @@ def create(self, label_dir, model_type, filename=None):
             if self.filename is not None:
                 filename = self.filename
             else:
-                filename="image"
+                filename = "image"
             labels = [str(filename) for i in range(self.image_dims[0])]
 
         df = pd.DataFrame(

src/napari_cellseg3d/plugin_model_training.py

@@ -32,6 +32,7 @@
 from napari_cellseg3d.model_workers import TrainingWorker
 
 NUMBER_TABS = 3
+DEFAULT_PATCH_SIZE = 60
 
 
 class Trainer(ModelFramework):
@@ -116,6 +117,7 @@ def __init__(
         self.data_path = ""
         self.label_path = ""
         self.results_path = ""
+        self.results_path_folder = ""
         ######################
         ######################
         ######################
@@ -252,7 +254,9 @@ def __init__(
         ]
         """Close buttons list for each tab"""
 
-        self.patch_size_widgets = ui.make_n_spinboxes(3, 10, 1023, 120)
+        self.patch_size_widgets = ui.make_n_spinboxes(
+            3, 10, 1024, DEFAULT_PATCH_SIZE
+        )
 
         self.patch_size_lbl = [
             QLabel(f"Size of patch in {axis} :") for axis in "xyz"
@@ -651,23 +655,24 @@ def build(self):
         ui.make_scrollable(
             contained_layout=data_tab_layout,
             containing_widget=data_tab,
-            min_wh=[100, 150],
+            min_wh=[200, 300],
         )  # , max_wh=[200,1000])
         self.addTab(data_tab, "Data")
 
         ui.make_scrollable(
             contained_layout=augment_tab_l,
             containing_widget=augment_tab_w,
-            min_wh=[100, 200],
+            min_wh=[200, 300],
         )
         self.addTab(augment_tab_w, "Augmentation")
 
         ui.make_scrollable(
             contained_layout=train_tab_layout,
             containing_widget=train_tab,
-            min_wh=[100, 200],
+            min_wh=[200, 300],
         )
         self.addTab(train_tab, "Training")
+        self.setMinimumSize(220, 200)
 
     def show_dialog_lab(self):
         """Shows the  dialog to load label files in a path, loads them (see :doc:model_framework) and changes the widget
@@ -760,12 +765,12 @@ def start(self):
                 "class": self.get_model(self.model_choice.currentText()),
                 "name": self.model_choice.currentText(),
             }
-            self.results_path = (
+            self.results_path_folder = (
                 self.results_path
                 + f"/{model_dict['name']}_results_{utils.get_date_time()}"
             )
             os.makedirs(
-                self.results_path, exist_ok=False
+                self.results_path_folder, exist_ok=False
             )  # avoid overwrite where possible
 
             if self.use_transfer_choice.isChecked():
@@ -777,7 +782,7 @@ def start(self):
                 weights_path = None
 
             self.log.print_and_log(
-                f"Notice : Saving results to : {self.results_path}"
+                f"Notice : Saving results to : {self.results_path_folder}"
             )
 
             self.worker = TrainingWorker(
@@ -790,7 +795,7 @@ def start(self):
                 learning_rate=self.learning_rate,
                 val_interval=self.val_interval,
                 batch_size=self.batch_size,
-                results_path=self.results_path,
+                results_path=self.results_path_folder,
                 sampling=self.patch_choice.isChecked(),
                 num_samples=self.num_samples,
                 sample_size=self.patch_size,
@@ -812,11 +817,13 @@ def start(self):
             self.worker.errored.connect(self.on_error)
 
         if self.worker.is_running:
+            self.log.print_and_log("*" * 20)
             self.log.print_and_log(
-                f"Stop requested at {utils.get_time()}. \nWaiting for next validation step..."
+                f"Stop requested at {utils.get_time()}. \nWaiting for next yielding step..."
             )
             self.stop_requested = True
-            self.btn_start.setText("Stopping... Please wait for next saving")
+            self.btn_start.setText("Stopping... Please wait")
+            self.log.print_and_log("*" * 20)
             self.worker.quit()
         else:
             self.worker.start()
@@ -833,32 +840,35 @@ def on_start(self):
 
     def on_finish(self):
         """Catches finished signal from worker"""
+        self.log.print_and_log("*" * 20)
         self.log.print_and_log(f"\nWorker finished at {utils.get_time()}")
 
-        self.log.print_and_log(f"Saving last loss plot at {self.results_path}")
+        self.log.print_and_log(f"Saving in {self.results_path_folder}")
+        self.log.print_and_log(f"Saving last loss plot")
 
         if self.canvas is not None:
             self.canvas.figure.savefig(
                 (
-                    self.results_path
+                    self.results_path_folder
                     + f"/final_metric_plots_{utils.get_time_filepath()}.png"
                 ),
                 format="png",
             )
 
-        self.log.print_and_log("Auto-saving log")
-        self.save_log()
+        self.log.print_and_log("Saving log")
+        self.save_log(spec_path=self.results_path_folder)
 
         self.log.print_and_log("Done")
         self.log.print_and_log("*" * 10)
 
         self.btn_start.setText("Start")
         [btn.setVisible(True) for btn in self.close_buttons]
 
+        del self.worker
         self.worker = None
+        self.results_path_folder = ""
         self.empty_cuda_cache()
 
-        self.results_path = ""
         # self.clean_cache() # trying to fix memory leak
 
     def on_error(self):
@@ -880,13 +890,17 @@ def on_yield(data, widget):
             widget.update_loss_plot(data["losses"], data["val_metrics"])
 
         if widget.stop_requested:
+            widget.log.print_and_log(
+                "Saving weights from aborted training in results folder"
+            )
             torch.save(
                 data["weights"],
                 os.path.join(
-                    widget.results_path,
+                    widget.results_path_folder,
                     f"latest_weights_aborted_training_{utils.get_time_filepath()}.pth",
                 ),
             )
+            widget.log.print_and_log("Saving complete")
             widget.stop_requested = False
 
     # def clean_cache(self):
@@ -942,7 +956,7 @@ def plot_loss(self, loss, dice_metric):
             )
             self.canvas.draw_idle()
 
-            plot_path = self.results_path + "/Loss_plots"
+            plot_path = self.results_path_folder + "/Loss_plots"
             os.makedirs(plot_path, exist_ok=True)
             if self.canvas is not None:
                 self.canvas.figure.savefig(