test: adjust tests

Author: okunator

cellseg_models_pytorch/datasets/tests/__init__.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from cellseg_models_pytorch.decoders import UnetDecoder
+from cellseg_models_pytorch.decoders.unet_decoder import UnetDecoder
 
 
 @pytest.mark.parametrize(

cellseg_models_pytorch/decoders/tests/test_decoders.py

@@ -1,51 +1,52 @@
-import pytest
-
-from cellseg_models_pytorch.inference import ResizeInferer, SlidingWindowInferer
-from cellseg_models_pytorch.models import cellpose_plus
-
-
-@pytest.mark.parametrize("batch_size", [1, 2])
-def test_slidingwin_inference(img_dir, batch_size):
-    model = cellpose_plus(sem_classes=3, type_classes=3, long_skip="unet")
-
-    inferer = SlidingWindowInferer(
-        model,
-        img_dir,
-        out_activations={"sem": "softmax", "type": "softmax", "cellpose": "tanh"},
-        out_boundary_weights={"sem": False, "type": False, "cellpose": True},
-        patch_size=(256, 256),
-        stride=256,
-        padding=80,
-        instance_postproc="hovernet",
-        batch_size=batch_size,
-        device="cpu",
-        parallel=False,
-    )
-
-    inferer.infer()
-
-    samples = list(inferer.out_masks.keys())
-    assert inferer.out_masks[samples[0]]["inst"].shape == (512, 512)
-
-
-@pytest.mark.parametrize("batch_size", [1, 2])
-def test_resize_inference(img_dir, batch_size):
-    model = cellpose_plus(sem_classes=3, type_classes=3, long_skip="unet")
-
-    inferer = ResizeInferer(
-        model,
-        img_dir,
-        out_activations={"sem": "softmax", "type": "softmax", "cellpose": "tanh"},
-        out_boundary_weights={"sem": False, "type": False, "cellpose": True},
-        resize=(256, 256),
-        padding=80,
-        instance_postproc="hovernet",
-        batch_size=batch_size,
-        device="cpu",
-        parallel=False,
-    )
-
-    inferer.infer()
-
-    samples = list(inferer.out_masks.keys())
-    assert inferer.out_masks[samples[0]]["inst"].shape == (512, 512)
+# import pytest
+
+# from cellseg_models_pytorch.inference import Inferer, SlidingWindowInferer
+# from cellseg_models_pytorch.models import cellpose_plus
+# import numpy as np
+# import torch
+
+# def test_slidingwin_inferer(img_sample512):
+#     model = cellpose_plus(n_sem_classes=3, n_type_classes=3, long_skip="unet")
+
+#     inferer = SlidingWindowInferer(
+#         model,
+#         patch_shape=(256, 256),
+#         stride=256,
+#         out_activations={"sem-sem": "softmax", "cellpose-type": "softmax", "cellpose-cellpose": "tanh"},
+#         out_boundary_weights={"sem-sem": False, "cellpose-type": False, "cellpose-cellpose": True},
+#         post_proc_method="cellpose",
+#         num_post_proc_threads=1,
+#         mixed_precision=True,
+#     )
+
+#     im = img_sample512
+#     im = np.transpose(im, (2, 0, 1))  
+#     im = np.expand_dims(im, axis=0)  
+#     im = torch.tensor(im)
+#     probs = inferer.predict(im.float())
+#     out_masks = inferer.post_process(probs)
+
+#     assert out_masks["sem-sem"].shape == (256, 256)
+
+
+# def test_inferer(img_sample256):
+#     model = cellpose_plus(n_sem_classes=3, n_type_classes=3, long_skip="unet")
+
+#     inferer = Inferer(
+#         model,
+#         input_shape=(256, 256),
+#         out_activations={"sem-sem": "softmax", "cellpose-type": "softmax", "cellpose-cellpose": "tanh"},
+#         out_boundary_weights={"sem-sem": False, "cellpose-type": False, "cellpose-cellpose": True},
+#         post_proc_method="cellpose",
+#         num_post_proc_threads=1,
+#         mixed_precision=True,
+#     )
+
+#     im = img_sample256
+#     im = np.transpose(im, (2, 0, 1))  
+#     im = np.expand_dims(im, axis=0)  
+#     im = torch.tensor(im)
+#     probs = inferer.predict(im.float())
+#     out_masks = inferer.post_process(probs)
+
+#     assert out_masks["sem-sem"].shape == (256, 256)

cellseg_models_pytorch/inference/tests/test_inference.py

@@ -1,5 +1,3 @@
-from collections import OrderedDict
-
 import pytest
 import torch
 
@@ -20,16 +18,14 @@
     "losses",
     [
         {"inst": JointLoss([TverskyLoss(), IoULoss()])},
-        {"inst": CELoss(), "type": SSIM()},
+        {"inst": CELoss(), "types": SSIM()},
     ],
 )
 def test_multitask_loss(n_classes, losses):
-    losses = OrderedDict(losses)
-
     yhats = {}
     targets = {}
     for i, br in enumerate(losses.keys(), 1):
-        yhats[f"{br}"], targets[f"{br}"] = _get_dummy_pair(n_classes + i)
+        yhats[br], targets[br] = _get_dummy_pair(n_classes + i)
 
     mtl = MultiTaskLoss(losses)
     mtl(yhats, targets)

cellseg_models_pytorch/losses/tests/test_multitask_loss.py

@@ -1,36 +1,32 @@
 import pytest
 import torch
 
-from cellseg_models_pytorch.models import MultiTaskUnet, get_model
+from cellseg_models_pytorch.models import get_model
 
 
 @pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
-@pytest.mark.parametrize("style_channels", [None, 32])
-@pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_cppnet_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
+def test_cppnet_fwdbwd(enc_name, model_type):
     n_rays = 3
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="cppnet",
         type=model_type,
         enc_name=enc_name,
         n_rays=n_rays,
-        ntypes=3,
-        ntissues=3,
-        style_channels=style_channels,
-        add_stem_skip=add_stem_skip,
+        n_type_classes=3,
+        n_sem_classes=3,
         enc_pretrain=False,
     )
 
     y = model(x)
-    y["stardist_refined"].mean().backward()
+    y["stardist-stardist"].mean().backward()
 
-    assert y["type"].shape == x.shape
-    assert y["stardist_refined"].shape == torch.Size([1, n_rays, 64, 64])
+    assert y["type-type"].shape == x.shape
+    assert y["stardist-stardist"].shape == torch.Size([1, n_rays, 64, 64])
 
-    if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
+    if "sem-sem" in y.keys():
+        assert y["sem-sem"].shape == torch.Size([1, 3, 64, 64])
 
 
 @pytest.mark.parametrize(
@@ -43,151 +39,119 @@ def test_cppnet_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
     ],
 )
 @pytest.mark.parametrize("model_type", ["base", "plus", "small_plus", "small"])
-@pytest.mark.parametrize("style_channels", [None, 32])
-def test_cellvit_fwdbwd(enc_name, model_type, style_channels):
+def test_cellvit_fwdbwd(enc_name, model_type):
     x = torch.rand([1, 3, 32, 32])
     model = get_model(
         name="cellvit",
         type=model_type,
         enc_name=enc_name,
-        ntypes=3,
-        ntissues=3,
-        style_channels=style_channels,
+        n_type_classes=3,
+        n_sem_classes=3,
         enc_pretrain=False,
+        enc_freeze=True
     )
-    model.freeze_encoder()
 
     y = model(x)
-    y["hovernet"].mean().backward()
+    y["hovernet-hovernet"].mean().backward()
 
-    assert y["type"].shape == x.shape
+    assert y["type-type"].shape == x.shape
 
-    if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 32, 32])
+    if "sem-sem" in y.keys():
+        assert y["sem-sem"].shape == torch.Size([1, 3, 32, 32])
 
 
 @pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus", "small_plus", "small"])
-@pytest.mark.parametrize("style_channels", [None, 32])
-@pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_hovernet_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
+@pytest.mark.parametrize("stem_skip_kws", [None, {"short_skip": "residual"}])
+@pytest.mark.parametrize("style_channels", [None, 256])
+def test_hovernet_fwdbwd(enc_name, model_type, stem_skip_kws, style_channels):
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="hovernet",
         type=model_type,
         enc_name=enc_name,
-        ntypes=3,
-        ntissues=3,
-        style_channels=style_channels,
-        add_stem_skip=add_stem_skip,
+        n_type_classes=3,
+        n_sem_classes=3,
         enc_pretrain=False,
+        style_channels=style_channels,
+        stem_skip_kws=stem_skip_kws,
     )
 
     y = model(x)
-    y["hovernet"].mean().backward()
+    y["hovernet-hovernet"].mean().backward()
 
-    assert y["type"].shape == x.shape
+    assert y["type-type"].shape == x.shape
 
-    if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
+    if "sem-sem" in y.keys():
+        assert y["sem-sem"].shape == torch.Size([1, 3, 64, 64])
 
 
 @pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
-@pytest.mark.parametrize("style_channels", [None, 32])
-@pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_stardist_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
+def test_stardist_fwdbwd(enc_name, model_type):
     n_rays = 3
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="stardist",
         type=model_type,
         n_rays=n_rays,
         enc_name=enc_name,
-        ntypes=3,
-        ntissues=3,
-        style_channels=style_channels,
-        add_stem_skip=add_stem_skip,
+        n_type_classes=3,
+        n_sem_classes=3,
         enc_pretrain=False,
     )
 
     y = model(x)
-    y["stardist"].mean().backward()
+    y["stardist-stardist"].mean().backward()
 
-    assert y["type"].shape == x.shape
-    assert y["stardist"].shape == torch.Size([1, n_rays, 64, 64])
+    assert y["stardist-type"].shape == x.shape
+    assert y["stardist-stardist"].shape == torch.Size([1, n_rays, 64, 64])
 
-    if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
+    if "sem-sem" in y.keys():
+        assert y["sem-sem"].shape == torch.Size([1, 3, 64, 64])
 
 
 @pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
-@pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_cellpose_fwdbwd(enc_name, model_type, add_stem_skip):
+def test_cellpose_fwdbwd(enc_name, model_type):
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="cellpose",
         type=model_type,
         enc_name=enc_name,
-        ntypes=3,
-        ntissues=3,
-        add_stem_skip=add_stem_skip,
+        n_type_classes=3,
+        n_sem_classes=3,
         enc_pretrain=False,
     )
 
     y = model(x)
-    y["cellpose"].mean().backward()
+    y["cellpose-cellpose"].mean().backward()
 
-    assert y["type"].shape == x.shape
-    assert y["cellpose"].shape == torch.Size([1, 2, 64, 64])
+    assert y["cellpose-type"].shape == x.shape
+    assert y["cellpose-cellpose"].shape == torch.Size([1, 2, 64, 64])
 
-    if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
+    if "sem-sem" in y.keys():
+        assert y["sem-sem"].shape == torch.Size([1, 3, 64, 64])
 
 
 @pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
-@pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_cellpose_fwdbwd(enc_name, model_type, add_stem_skip):
+def test_cellpose_fwdbwd(enc_name, model_type):
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="omnipose",
         type=model_type,
         enc_name=enc_name,
-        ntypes=3,
-        ntissues=3,
-        add_stem_skip=add_stem_skip,
+        n_type_classes=3,
+        n_sem_classes=3,
         enc_pretrain=False,
     )
 
     y = model(x)
-    y["omnipose"].mean().backward()
-
-    assert y["type"].shape == x.shape
-    assert y["omnipose"].shape == torch.Size([1, 2, 64, 64])
-
-    if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
+    y["omnipose-omnipose"].mean().backward()
 
+    assert y["omnipose-type"].shape == x.shape
+    assert y["omnipose-omnipose"].shape == torch.Size([1, 2, 64, 64])
 
-@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
-@pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_multitaskunet_fwdbwd(enc_name, add_stem_skip):
-    x = torch.rand([1, 3, 64, 64])
-    m = MultiTaskUnet(
-        decoders=("sem",),
-        heads={"sem": {"sem": 3}},
-        n_conv_layers={"sem": (1, 1, 1, 1)},
-        n_conv_blocks={"sem": ((2,), (2,), (2,), (2,))},
-        out_channels={"sem": (128, 64, 32, 16)},
-        long_skips={"sem": "unet"},
-        dec_params={"sem": None},
-        add_stem_skip=add_stem_skip,
-        enc_name=enc_name,
-        enc_pretrain=False,
-    )
-    y = m(x)
-    y["sem"].mean().backward()
-
-    assert y["sem"].shape == torch.Size([1, 3, 64, 64])
+    if "sem-sem" in y.keys():
+        assert y["sem-sem"].shape == torch.Size([1, 3, 64, 64])