X_SSL_Net/tests/test_xnet_2d.py

from __future__ import annotations

import importlib
import sys
import warnings
from pathlib import Path

import torch
from torch import nn

ROOT_DIR = Path(__file__).resolve().parents[1]
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))


def test_importing_xnet2d_does_not_emit_deprecation_warnings() -> None:
    modules_to_clear = [
        name
        for name in sys.modules
        if name == "lib.modules.xnet_2d" or name.startswith("lib.modules.lib_mamba")
    ]
    for name in modules_to_clear:
        sys.modules.pop(name, None)

    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always", DeprecationWarning)
        importlib.import_module("lib.modules.xnet_2d")

    assert not [
        warning
        for warning in caught
        if issubclass(warning.category, DeprecationWarning)
    ]


def test_xnet2d_forward_preserves_segmentation_shape() -> None:
    from lib.modules.xnet_2d import XNet2d, XTEB2d

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = XNet2d(
        in_channels=3,
        num_classes=1,
        encoder_channels=(8, 16, 24, 32),
        encoder_depths=(1, 1, 1, 1),
        decoder_channels=(24, 16, 8),
        stem_channels=8,
        bottleneck_depth=1,
        global_ratio=1.0,
        use_wavelet_branch=True,
        use_global_branch_stage1=False,
        ssm_d_state=1,
        ssm_backend="torch",
        use_frequency_refine=True,
        learnable_low_freq_radius=False,
    ).to(device)
    if device.type == "cpu":
        for module in model.modules():
            if isinstance(module, XTEB2d):
                module.global_branch = nn.Identity()
    model.eval()

    x = torch.randn(2, 3, 64, 64, device=device)
    with torch.no_grad():
        outputs = model(x)

    assert outputs["seg_logits"].shape == (2, 1, 64, 64)
    assert outputs["logits"].shape == outputs["seg_logits"].shape
    assert "guides" not in outputs


def test_xnet2d_decoder_uses_plain_unet_skip_connections() -> None:
    from lib.modules.xnet_2d import XNet2d

    model = XNet2d(
        in_channels=3,
        num_classes=1,
        encoder_channels=(8, 16, 24, 32),
        encoder_depths=(1, 1, 1, 1),
        decoder_channels=(24, 16, 8),
        stem_channels=8,
        bottleneck_depth=1,
        use_global_branch_stage1=False,
        ssm_d_state=1,
        ssm_backend="torch",
    )

    decoder_module_names = dict(model.decoder.named_modules())

    assert not any(name.startswith("guide") for name in decoder_module_names)


def test_xnet2d_main_path_does_not_use_batchnorm2d() -> None:
    from lib.modules.xnet_2d import XNet2d

    model = XNet2d(
        in_channels=3,
        num_classes=1,
        encoder_channels=(8, 16, 24, 32),
        encoder_depths=(1, 1, 1, 1),
        decoder_channels=(24, 16, 8),
        stem_channels=8,
        bottleneck_depth=1,
        use_global_branch_stage1=False,
        ssm_d_state=1,
        ssm_backend="torch",
    )

    batchnorm_modules = [
        module
        for module in model.modules()
        if isinstance(module, nn.BatchNorm2d)
    ]

    assert not batchnorm_modules


def test_xnet2d_train_mode_supports_batch_size_one_without_batchnorm_failure() -> None:
    from lib.modules.xnet_2d import XNet2d, XTEB2d

    model = XNet2d(
        in_channels=3,
        num_classes=1,
        encoder_channels=(8, 16, 24, 32),
        encoder_depths=(1, 1, 1, 1),
        decoder_channels=(24, 16, 8),
        stem_channels=8,
        bottleneck_depth=1,
        global_ratio=1.0,
        use_wavelet_branch=True,
        use_global_branch_stage1=False,
        ssm_d_state=1,
        ssm_backend="torch",
        use_frequency_refine=True,
        learnable_low_freq_radius=False,
    )
    for module in model.modules():
        if isinstance(module, XTEB2d):
            module.global_branch = nn.Identity()
    model.train()

    x = torch.randn(1, 3, 64, 64)
    outputs = model(x)

    assert outputs["seg_logits"].shape == (1, 1, 64, 64)