quadra 0.0.1__py3-none-any.whl → 2.1.13__py3-none-any.whl

quadra/utils/tests/fixtures/models/anomaly.py

@@ -0,0 +1,89 @@
+import pytest
+import torch
+from anomalib.models.draem.torch_model import DraemModel
+from anomalib.models.efficient_ad.torch_model import EfficientAdModel
+from anomalib.models.padim.torch_model import PadimModel
+from anomalib.models.patchcore.torch_model import PatchcoreModel
+
+
+@pytest.fixture
+def padim_resnet18():
+    """Yield a padim model with resnet18 encoder."""
+    yield PadimModel(
+        input_size=[224, 224],  # TODO: This is hardcoded may be not a good idea
+        backbone="resnet18",
+        layers=["layer1", "layer2", "layer3"],
+        pretrained_weights=None,
+        tied_covariance=False,
+        pre_trained=False,
+    )
+
+
+@torch.inference_mode()
+def _initialize_patchcore_model(patchcore_model: PatchcoreModel, coreset_sampling_ratio: float = 0.1) -> PatchcoreModel:
+    """Initialize a Patchcore model by simulating a training step.
+
+    Args:
+        patchcore_model: Patchcore model to initialize
+        coreset_sampling_ratio: Coreset sampling ratio to use for the initialization
+
+    Returns:
+        Patchcore model with initialized memory bank
+    """
+    with torch.no_grad():
+        training_features = None
+        random_input = torch.rand([1, 3, *patchcore_model.input_size])
+
+        if training_features is None:
+            training_features = patchcore_model(random_input)
+        else:
+            training_features = torch.cat([training_features, patchcore_model(random_input)], dim=0)
+
+        patchcore_model.eval()
+        patchcore_model.subsample_embedding(training_features, sampling_ratio=coreset_sampling_ratio)
+
+        # Simulate a memory bank with 5 images, at the current stage patchcore onnx export is not handling
+        # large memory banks well, so we are using a small one for the benchmark
+        memory_bank_number, memory_bank_n_features = patchcore_model.memory_bank.shape
+        patchcore_model.memory_bank = torch.rand([5 * memory_bank_number, memory_bank_n_features])
+        patchcore_model.train()
+
+    return patchcore_model
+
+
+@pytest.fixture
+def patchcore_resnet18():
+    """Yield a patchcore model with resnet18 encoder."""
+    model = PatchcoreModel(
+        input_size=[224, 224],  # TODO: This is hardcoded may be not a good idea
+        backbone="resnet18",
+        layers=["layer2", "layer3"],
+        pre_trained=False,
+    )
+
+    yield _initialize_patchcore_model(model)
+
+
+@pytest.fixture
+def draem():
+    """Yield a draem model."""
+    yield DraemModel()
+
+
+@pytest.fixture
+def efficient_ad_small():
+    """Yield a draem model."""
+
+    class EfficientAdForwardWrapper(EfficientAdModel):
+        """Wrap the forward method to avoid passing optional parameters."""
+
+        def forward(self, x):
+            return super().forward(x, None)
+
+    model = EfficientAdForwardWrapper(
+        teacher_out_channels=384,
+        input_size=[256, 256],  # TODO: This is hardcoded may be not a good idea
+        pretrained_teacher_type="nelson",
+    )
+
+    yield model

quadra/utils/tests/fixtures/models/classification.py

@@ -0,0 +1,45 @@
+import pytest
+
+from quadra.models.classification import TimmNetworkBuilder, TorchHubNetworkBuilder
+
+
+@pytest.fixture
+def resnet18():
+    """Yield a resnet18 model."""
+    yield TimmNetworkBuilder("resnet18", pretrained=False, freeze=True, exportable=True)
+
+
+@pytest.fixture
+def resnet50():
+    """Yield a resnet50 model."""
+    yield TimmNetworkBuilder("resnet50", pretrained=False, freeze=True, exportable=True)
+
+
+@pytest.fixture
+def vit_tiny_patch16_224():
+    """Yield a vit_tiny_patch16_224 model."""
+    yield TimmNetworkBuilder("vit_tiny_patch16_224", pretrained=False, freeze=True, exportable=True)
+
+
+@pytest.fixture
+def dino_vits8():
+    """Yield a dino_vits8 model."""
+    yield TorchHubNetworkBuilder(
+        repo_or_dir="facebookresearch/dino:main",
+        model_name="dino_vits8",
+        pretrained=False,
+        freeze=True,
+        exportable=True,
+    )
+
+
+@pytest.fixture
+def dino_vitb8():
+    """Yield a dino_vitb8 model."""
+    yield TorchHubNetworkBuilder(
+        repo_or_dir="facebookresearch/dino:main",
+        model_name="dino_vitb8",
+        pretrained=False,
+        freeze=True,
+        exportable=True,
+    )

quadra/utils/tests/fixtures/models/segmentation.py

@@ -0,0 +1,33 @@
+import pytest
+
+from quadra.modules.backbone import create_smp_backbone
+
+
+@pytest.fixture
+def smp_resnet18_unet():
+    """Yield a unet with resnet18 encoder."""
+    yield create_smp_backbone(
+        arch="unet",
+        encoder_name="resnet18",
+        encoder_weights=None,
+        encoder_depth=5,
+        freeze_encoder=True,
+        in_channels=3,
+        num_classes=1,
+        activation=None,
+    )
+
+
+@pytest.fixture
+def smp_resnet18_unetplusplus():
+    """Yield a unetplusplus with resnet18 encoder."""
+    yield create_smp_backbone(
+        arch="unetplusplus",
+        encoder_name="resnet18",
+        encoder_weights=None,
+        encoder_depth=5,
+        freeze_encoder=True,
+        in_channels=3,
+        num_classes=1,
+        activation=None,
+    )

quadra/utils/tests/helpers.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from hydra import compose, initialize_config_module
+from hydra.core.hydra_config import HydraConfig
+
+from quadra.main import main
+from quadra.utils.export import get_export_extension
+
+
+# taken from hydra unit tests
+def _random_image(size: tuple[int, int] = (10, 10)) -> np.ndarray:
+    """Generate random image."""
+    return np.random.randint(0, 255, size=size, dtype=np.uint8)
+
+
+def execute_quadra_experiment(overrides: list[str], experiment_path: Path) -> None:
+    """Execute quadra experiment."""
+    with initialize_config_module(config_module="quadra.configs", version_base="1.3.0"):
+        if not experiment_path.exists():
+            experiment_path.mkdir(parents=True)
+        os.chdir(experiment_path)
+        # cfg = compose(config_name="config", overrides=overrides)
+        cfg = compose(config_name="config", overrides=overrides, return_hydra_config=True)
+        # workaround without actual main function
+        # check https://github.com/facebookresearch/hydra/issues/2017 for more details
+        HydraConfig.instance().set_config(cfg)
+
+        main(cfg)
+
+
+def check_deployment_model(export_type: str):
+    """Check that the runtime model is present and valid.
+
+    Args:
+        export_type: The type of the exported model.
+    """
+    extension = get_export_extension(export_type)
+
+    assert os.path.exists(f"deployment_model/model.{extension}")
+    assert os.path.exists("deployment_model/model.json")
+
+
+def get_quadra_test_device():
+    """Get the device to use for the tests. If the QUADRA_TEST_DEVICE environment variable is set, it is used."""
+    return os.environ.get("QUADRA_TEST_DEVICE", "cpu")
+
+
+def setup_trainer_for_lightning() -> list[str]:
+    """Setup trainer for lightning depending on the device. If cuda is used, the device index is also set.
+    If cpu is used, the trainer is set to lightning_cpu.
+
+    Returns:
+        A list of overrides for the trainer.
+    """
+    overrides = []
+    device = get_quadra_test_device()
+    torch_device = torch.device(device)
+    if torch_device.type == "cuda":
+        device_index = torch_device.index
+        overrides.append("trainer=lightning_gpu")
+        overrides.append(f"trainer.devices=[{device_index}]")
+    else:
+        overrides.append("trainer=lightning_cpu")
+
+    return overrides

quadra/utils/tests/models.py

@@ -0,0 +1,27 @@
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class SingleInputModel(nn.Module):
+    """Model taking a single input."""
+
+    def forward(self, x: Any):
+        return x
+
+
+class DoubleInputModel(nn.Module):
+    """Model taking two inputs."""
+
+    def forward(self, x: Any, y: Any):
+        return x, y
+
+
+class UnsupportedInputModel(nn.Module):
+    """Model taking an unsupported input."""
+
+    def forward(self, x: torch.Tensor, y: str):
+        y = f"unsupported input: {y}"
+
+        return x