Improving testing

Author: C-Achard

napari_cellseg3d/_tests/test_plugin_inference.py

@@ -1,9 +1,10 @@
 from tifffile import imread
 from pathlib import Path
 
+from napari_cellseg3d.config import MODEL_LIST
 from napari_cellseg3d._tests.fixtures import LogFixture
 from napari_cellseg3d.code_plugins.plugin_model_inference import Inferer
-
+from napari_cellseg3d.code_models.models.model_test import TestModel
 
 def test_inference(make_napari_viewer, qtbot):
 
@@ -26,10 +27,14 @@ def test_inference(make_napari_viewer, qtbot):
 
     assert widget.check_ready()
 
-    # widget.start()  # takes too long on Github Actions
-    # assert widget.worker is not None
+    MODEL_LIST["test"] = TestModel
+    widget.model_choice.addItem("test")
+    widget.setCurrentIndex(-1)
+
+    widget.start()  # takes too long on Github Actions
+    assert widget.worker is not None
 
-    # with qtbot.waitSignal(signal=widget.worker.yielded, timeout=60000, raising=False) as blocker:
-    #     blocker.connect(widget.worker.errored)
+    with qtbot.waitSignal(signal=widget.worker.finished, timeout=60000, raising=False) as blocker:
+        blocker.connect(widget.worker.errored)
 
     # assert len(viewer.layers) == 2

napari_cellseg3d/_tests/test_training.py

@@ -3,6 +3,8 @@
 from napari_cellseg3d import config
 from napari_cellseg3d.code_plugins.plugin_model_training import Trainer
 from napari_cellseg3d._tests.fixtures import LogFixture
+from napari_cellseg3d.config import MODEL_LIST
+from napari_cellseg3d.code_models.models.model_test import TestModel
 
 
 def test_training(make_napari_viewer, qtbot):
@@ -30,12 +32,15 @@ def test_training(make_napari_viewer, qtbot):
     #################
     # Training is too long to test properly this way. Do not use on Github
     #################
+    MODEL_LIST["test"] = TestModel()
+    widget.model_choice.addItem("test")
+    widget.model_choice.setCurrentIndex(len(MODEL_LIST.keys())-1)
 
-    # widget.start()
-    # assert widget.worker is not None
+    widget.start()
+    assert widget.worker is not None
 
-    # with qtbot.waitSignal(signal=widget.worker.yielded, timeout=60000, raising=False) as blocker:  # wait only for 60 seconds.
-    #     blocker.connect(widget.worker.errored)
+    with qtbot.waitSignal(signal=widget.worker.finished, timeout=10000, raising=False) as blocker:  # wait only for 60 seconds.
+        blocker.connect(widget.worker.errored)
 
 
 def test_update_loss_plot(make_napari_viewer):

napari_cellseg3d/code_models/model_workers.py

@@ -455,7 +455,7 @@ def model_output(
         inputs = inputs.to("cpu")
 
         model_output = lambda inputs: post_process_transforms(
-            self.config.model_info.get_model().get_output(model, inputs)
+            self.config.model_info.get_model().get_output(model, inputs)  # TODO(cyril) refactor those functions
         )
 
         def model_output(inputs):
@@ -870,7 +870,7 @@ def inference(self):
             model.to("cpu")
 
         except Exception as e:
-            self.log(f"Error : {e}")
+            self.log(f"Error during inference : {e}")
             self.quit()
         finally:
             self.quit()
@@ -1078,6 +1078,7 @@ def train(self):
                 torch.set_num_threads(1)
                 self.log("Number of threads has been set to 1 for macOS")
 
+            self.log(f"config model : {self.config.model_info.name}")
             model_name = model_config.name
             model_class = model_config.get_model()
 
@@ -1314,7 +1315,7 @@ def train(self):
                         )
                     )
                 except RuntimeError as e:
-                    logger.error(f"Error : {e}")
+                    logger.error(f"Error when loading weights : {e}")
                     warn = (
                         "WARNING:\nIt'd seem that the weights were incompatible with the model,\n"
                         "the model will be trained from random weights"
@@ -1333,6 +1334,9 @@ def train(self):
 
             device = self.config.device
 
+            if model_name == "test":
+                self.quit()
+
             for epoch in range(self.config.max_epochs):
                 # self.log("\n")
                 self.log("-" * 10)
@@ -1472,7 +1476,7 @@ def train(self):
             model.to("cpu")
 
         except Exception as e:
-            self.log(f"Error : {e}")
+            self.log(f"Error in training : {e}")
             self.quit()
         finally:
             self.quit()

napari_cellseg3d/code_models/models/model_test.py

@@ -0,0 +1,36 @@
+import torch
+from torch import nn
+
+
+def get_weights_file():
+    return "test.pth"
+
+
+class TestModel(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(1, 1)
+
+    def forward(self, x):
+        return self.linear(x)
+
+    def get_net(self):
+        return self
+
+    def get_output(self, _, input):
+        return input
+
+    def get_validation(self, val_inputs):
+        return val_inputs
+
+if __name__ == "__main__":
+
+    model = TestModel()
+    model.train()
+    model.zero_grad()
+    from napari_cellseg3d.config import WEIGHTS_DIR
+    torch.save(
+        model.state_dict(),
+        WEIGHTS_DIR + f"/{get_weights_file()}"
+    )

napari_cellseg3d/code_plugins/plugin_model_training.py

@@ -791,6 +791,7 @@ def start(self):
                 self.data = None
                 raise err
 
+
             model_config = config.ModelInfo(
                 name=self.model_choice.currentText()
             )
@@ -820,6 +821,7 @@ def start(self):
 
             patch_size = [w.value() for w in self.patch_size_widgets]
 
+            logger.debug("Loading config...")
             self.worker_config = config.TrainingWorkerConfig(
                 device=self.get_device(),
                 model_info=model_config,

napari_cellseg3d/config.py

@@ -18,6 +18,9 @@
     model_TRAILMAP_MS as TRAILMAP_MS,
 )
 from napari_cellseg3d.code_models.models import model_VNet as VNet
+from napari_cellseg3d.utils import LOGGER
+
+logger = LOGGER
 
 # TODO(cyril) DOCUMENT !!! and add default values
 # TODO(cyril) add JSON load/save
@@ -28,6 +31,7 @@
     # "TRAILMAP": TRAILMAP,
     "TRAILMAP_MS": TRAILMAP_MS,
     "SwinUNetR": SwinUNetR,
+    # "test" : DO NOT USE, reserved for testing
 }
 
 INSTANCE_SEGMENTATION_METHOD_LIST = {
@@ -81,12 +85,14 @@ def get_model(self):
         try:
             return MODEL_LIST[self.name]
         except KeyError as e:
-            warnings.warn(f"Model {self.name} is not defined")
+            msg = f"Model {self.name} is not defined"
+            warnings.warn(msg)
+            logger.warning(msg)
             raise KeyError(e)
 
     @staticmethod
     def get_model_name_list():
-        print(
+        logger.info(
             f"Model list :\n" + str(f"{name}\n" for name in MODEL_LIST.keys())
         )
         return MODEL_LIST.keys()