Added WNet training test

Author: C-Achard

napari_cellseg3d/_tests/test_wnet_training.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+from napari_cellseg3d.code_models.models.wnet import train_wnet as t
+
+def test_wnet_training():
+    config = t.Config()
+
+    config.batch_size = 1
+    config.num_epochs = 1
+
+    config.train_volume_directory = str(Path(__file__).resolve().parent / "res/wnet_test")
+    config.eval_volume_directory = config.train_volume_directory
+    config.save_every = 1
+    config.val_interval = 2  # skip validation
+    config.save_model_path = config.train_volume_directory + "/test.pth"
+
+    ncuts_loss, rec_loss, model = t.train(train_config=config)
+
+    assert ncuts_loss is not None
+    assert rec_loss is not None
+    assert model is not None

napari_cellseg3d/code_models/models/wnet/train_wnet.py

@@ -720,8 +720,13 @@ def get_dataset(config):
         volume_directory=config.train_volume_directory
     )
     train_files = [d.get("image") for d in train_files]
+    # logger.debug(f"train_files: {train_files}")
     volumes = tiff.imread(train_files).astype(np.float32)
     volume_shape = volumes.shape
+    # logger.debug(f"volume_shape: {volume_shape}")
+
+    if len(volume_shape) == 3:
+        volumes = np.expand_dims(volumes, axis=0)
 
     if config.normalize_input:
         volumes = np.array(