Additional tests (#42)

* Additional tests

* Fixed weight compatibility test

* Added newest weight to HF for new WNet

* Improving worker tests

* Fix style

* Improve stats display test coverage

* Disable direct thread worker test for GH actions

* Update tox.ini

Fix ONNX reqs

* Update .coveragerc

Disable dev scripts coverage

* More tests & fixes

Tests :
    - Base plugin
    - ONNX inference
    - Cropping
Fixes :
    - Fixed review image loading not working
    - Fixed scaling when running inference on a folder and showing originals

* More worker tests

* More extensive training tests

* Fix anisotropy calculation

Author: C-Achard

.coveragerc

@@ -4,4 +4,4 @@ exclude_lines =
 
 [run]
 omit =
-    napari_cellseg3d/setup.py, napari_cellseg3d/code_models/models/wnet/train_wnet.py
+    napari_cellseg3d/setup.py, napari_cellseg3d/code_models/models/wnet/train_wnet.py, napari_cellseg3d/code_models/models/wnet/model.py,napari_cellseg3d/code_models/models/TEMPLATE_model.py, napari_cellseg3d/code_models/models/unet/*, napari_cellseg3d/dev_scripts/*

.gitignore

@@ -11,6 +11,9 @@ __pycache__/
 *.tiff
 napari_cellseg3d/_tests/res/*.csv
 *.pth
+*.pt
+*.onnx
+*.tar.gz
 *.db
 
 # Distribution / packaging

docs/res/code/utils.rst

@@ -27,10 +27,6 @@ get_time_filepath
 **************************************
 .. autofunction:: napari_cellseg3d.utils::get_time_filepath
 
-save_stack
-**************************************
-.. autofunction:: napari_cellseg3d.utils::save_stack
-
 get_padding_dim
 **************************************
 .. autofunction:: napari_cellseg3d.utils::get_padding_dim

napari_cellseg3d/_tests/test_base_plugin.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+from napari_cellseg3d.code_plugins.plugin_base import (
+    BasePluginSingleImage,
+)
+
+
+def test_base_single_image(make_napari_viewer_proxy):
+    viewer = make_napari_viewer_proxy()
+    plugin = BasePluginSingleImage(viewer)
+
+    test_folder = Path(__file__).parent.resolve()
+    test_image = str(test_folder / "res/test.tif")
+
+    assert plugin._check_results_path(str(test_folder))
+    plugin.image_path = test_image
+    assert plugin._default_path[0] != test_image
+    plugin._update_default()
+    assert plugin._default_path[0] == test_image

napari_cellseg3d/_tests/test_dock_widget.py

@@ -5,10 +5,10 @@
 from napari_cellseg3d.code_plugins.plugin_review_dock import Datamanager
 
 
-def test_prepare(make_napari_viewer):
+def test_prepare(make_napari_viewer_proxy):
     path_image = str(Path(__file__).resolve().parent / "res/test.tif")
     image = imread(str(path_image))
-    viewer = make_napari_viewer()
+    viewer = make_napari_viewer_proxy()
     viewer.add_image(image)
     widget = Datamanager(viewer)
     viewer.window.add_dock_widget(widget)

napari_cellseg3d/_tests/test_helper.py

@@ -1,8 +1,8 @@
 from napari_cellseg3d.code_plugins.plugin_helper import Helper
 
 
-def test_helper(make_napari_viewer):
-    viewer = make_napari_viewer()
+def test_helper(make_napari_viewer_proxy):
+    viewer = make_napari_viewer_proxy()
     widget = Helper(viewer)
 
     dock = viewer.window.add_dock_widget(widget)

napari_cellseg3d/_tests/test_interface.py

@@ -15,6 +15,6 @@ def test_log(qtbot):
 
 
 def test_zoom_factor():
-    resolution = [10.0, 10.0, 5.0]
+    resolution = [5.0, 10.0, 5.0]
     zoom = AnisotropyWidgets.anisotropy_zoom_factor(resolution)
-    assert zoom == [1, 1, 0.5]
+    assert zoom == [1, 0.5, 1]

napari_cellseg3d/_tests/test_model_framework.py

@@ -1,14 +1,15 @@
 from pathlib import Path
 
 from napari_cellseg3d.code_models import model_framework
+from napari_cellseg3d.config import MODEL_LIST
 
 
 def pth(path):
     return str(Path(path))
 
 
-def test_update_default(make_napari_viewer):
-    view = make_napari_viewer()
+def test_update_default(make_napari_viewer_proxy):
+    view = make_napari_viewer_proxy()
     widget = model_framework.ModelFramework(view)
 
     widget.images_filepaths = []
@@ -38,8 +39,8 @@ def test_update_default(make_napari_viewer):
     ]
 
 
-def test_create_train_dataset_dict(make_napari_viewer):
-    view = make_napari_viewer()
+def test_create_train_dataset_dict(make_napari_viewer_proxy):
+    view = make_napari_viewer_proxy()
     widget = model_framework.ModelFramework(view)
 
     widget.images_filepaths = [str(f"{i}.tif") for i in range(3)]
@@ -52,3 +53,64 @@ def test_create_train_dataset_dict(make_napari_viewer):
     ]
 
     assert widget.create_train_dataset_dict() == expect
+
+
+def test_log(make_napari_viewer_proxy):
+    mock_test = "test"
+    framework = model_framework.ModelFramework(
+        viewer=make_napari_viewer_proxy()
+    )
+    framework.log.print_and_log(mock_test)
+    assert len(framework.log.toPlainText()) != 0
+    assert framework.log.toPlainText() == "\n" + mock_test
+
+    framework.results_path = str(Path(__file__).resolve().parent / "res")
+    framework.save_log(do_timestamp=False)
+    log_path = Path(__file__).resolve().parent / "res/Log_report.txt"
+    assert log_path.is_file()
+    with Path.open(log_path.resolve(), "r") as f:
+        assert f.read() == "\n" + mock_test
+
+    # remove log file
+    log_path.unlink(missing_ok=False)
+    log_path = Path(__file__).resolve().parent / "res/Log_report.txt"
+    framework.save_log_to_path(str(log_path.parent), do_timestamp=False)
+    assert log_path.is_file()
+    with Path.open(log_path.resolve(), "r") as f:
+        assert f.read() == "\n" + mock_test
+    log_path.unlink(missing_ok=False)
+
+
+def test_display_elements(make_napari_viewer_proxy):
+    framework = model_framework.ModelFramework(
+        viewer=make_napari_viewer_proxy()
+    )
+
+    framework.display_status_report()
+    framework.display_status_report()
+
+    framework.custom_weights_choice.setChecked(False)
+    framework._toggle_weights_path()
+    assert not framework.weights_filewidget.isVisible()
+
+
+def test_available_models_retrieval(make_napari_viewer_proxy):
+    framework = model_framework.ModelFramework(
+        viewer=make_napari_viewer_proxy()
+    )
+    assert framework.get_available_models() == MODEL_LIST
+
+
+def test_update_weights_path(make_napari_viewer_proxy):
+    framework = model_framework.ModelFramework(
+        viewer=make_napari_viewer_proxy()
+    )
+    assert (
+        framework._update_weights_path(framework._default_weights_folder)
+        is None
+    )
+    name = str(Path.home() / "test/weight.pth")
+    framework._update_weights_path([name])
+    assert framework.weights_config.path == name
+    assert framework.weights_filewidget.text_field.text() == name
+    assert framework._default_weights_folder == str(Path.home() / "test")

napari_cellseg3d/_tests/test_models.py

@@ -1,4 +1,7 @@
+from pathlib import Path
+
 import numpy as np
+import pytest
 import torch
 from numpy.random import PCG64, Generator
 
@@ -8,6 +11,7 @@
     crf_batch,
     crf_with_config,
 )
+from napari_cellseg3d.code_models.models.model_TRAILMAP_MS import TRAILMAP_MS_
 from napari_cellseg3d.code_models.models.wnet.soft_Ncuts import SoftNCutsLoss
 from napari_cellseg3d.config import MODEL_LIST, CRFConfig
 
@@ -51,6 +55,15 @@ def test_soft_ncuts_loss():
     assert isinstance(res, torch.Tensor)
     assert 0 <= res <= 1  # ASSUMES NUMBER OF CLASS IS 2, NOT CORRECT IF K>2
 
+    loss = SoftNCutsLoss(
+        data_shape=[dims, dims, dims],
+        device="cpu",
+        intensity_sigma=4,
+        spatial_sigma=4,
+        radius=None,  # test radius=None init
+    )
+    assert loss.radius == 5
+
 
 def test_crf_batch():
     dims = 8
@@ -95,3 +108,33 @@ def on_yield(result):
 
     result = next(crf._run_crf_job())
     on_yield(result)
+
+
+def test_pretrained_weights_compatibility():
+    from napari_cellseg3d.code_models.workers import WeightsDownloader
+    from napari_cellseg3d.config import MODEL_LIST, PRETRAINED_WEIGHTS_DIR
+
+    for model_name in MODEL_LIST:
+        file_name = MODEL_LIST[model_name].weights_file
+        WeightsDownloader().download_weights(model_name, file_name)
+        model = MODEL_LIST[model_name](input_img_size=[128, 128, 128])
+        try:
+            model.load_state_dict(
+                torch.load(
+                    str(Path(PRETRAINED_WEIGHTS_DIR) / file_name),
+                    map_location="cpu",
+                ),
+                strict=True,
+            )
+        except RuntimeError:
+            pytest.fail(f"Failed to load weights for {model_name}")
+
+
+def test_trailmap_init():
+    test = TRAILMAP_MS_(
+        input_img_size=[128, 128, 128],
+        in_channels=1,
+        out_channels=1,
+        dropout_prob=0.3,
+    )
+    assert isinstance(test, TRAILMAP_MS_)

napari_cellseg3d/_tests/test_plugin_inference.py

@@ -1,10 +1,12 @@
 from pathlib import Path
 
+from numpy.random import PCG64, Generator
 from tifffile import imread
 
 from napari_cellseg3d._tests.fixtures import LogFixture
 from napari_cellseg3d.code_models.instance_segmentation import (
     INSTANCE_SEGMENTATION_METHOD_LIST,
+    volume_stats,
 )
 from napari_cellseg3d.code_models.models.model_test import TestModel
 from napari_cellseg3d.code_plugins.plugin_model_inference import (
@@ -13,14 +15,16 @@
 )
 from napari_cellseg3d.config import MODEL_LIST
 
+rand_gen = Generator(PCG64(12345))
 
-def test_inference(make_napari_viewer, qtbot):
+
+def test_inference(make_napari_viewer_proxy, qtbot):
     im_path = str(Path(__file__).resolve().parent / "res/test.tif")
     image = imread(im_path)
 
     assert image.shape == (6, 6, 6)
 
-    viewer = make_napari_viewer()
+    viewer = make_napari_viewer_proxy()
     widget = Inferer(viewer)
     widget.log = LogFixture()
     viewer.window.add_dock_widget(widget)
@@ -48,8 +52,8 @@ def test_inference(make_napari_viewer, qtbot):
     assert widget.worker_config is not None
     assert widget.model_info is not None
     widget.window_infer_box.setChecked(False)
-    worker = widget._create_worker_from_config(widget.worker_config)
 
+    worker = widget._create_worker_from_config(widget.worker_config)
     assert worker.config is not None
     assert worker.config.model_info is not None
     worker.config.layer = viewer.layers[0].data
@@ -65,5 +69,31 @@ def test_inference(make_napari_viewer, qtbot):
     assert isinstance(res, InferenceResult)
     assert res.result.shape == (8, 8, 8)
     assert res.instance_labels.shape == (8, 8, 8)
-
     widget.on_yield(res)
+
+    mock_image = rand_gen.random(size=(10, 10, 10))
+    mock_labels = rand_gen.integers(0, 10, (10, 10, 10))
+    mock_results = InferenceResult(
+        image_id=0,
+        original=mock_image,
+        instance_labels=mock_labels,
+        crf_results=mock_image,
+        stats=[volume_stats(mock_labels)],
+        result=mock_image,
+        model_name="test",
+    )
+    num_layers = len(viewer.layers)
+    widget.worker_config.post_process_config.instance.enabled = True
+    widget._display_results(mock_results)
+    assert len(viewer.layers) == num_layers + 4
+
+    # assert widget.check_ready()
+    # widget._setup_worker()
+    # # widget.config.show_results = True
+    # with qtbot.waitSignal(widget.worker.yielded, timeout=10000) as blocker:
+    #     blocker.connect(
+    #         widget.worker.errored
+    #     )  # Can add other signals to blocker
+    #     widget.worker.start()
+
+    assert widget.on_finish()