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

quadra/models/evaluation.py

@@ -0,0 +1,322 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, cast
+
+import numpy as np
+import torch
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from torch.jit import RecursiveScriptModule
+
+from quadra.utils.logger import get_logger
+
+try:
+    import onnxruntime as ort  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
+
+
+log = get_logger(__name__)
+
+
+class BaseEvaluationModel(ABC):
+    """Base interface for all evaluation models."""
+
+    def __init__(self, config: DictConfig) -> None:
+        self.model: Any
+        self.model_path: str | None
+        self.device: str
+        self.config = config
+        self.is_loaded = False
+        self.model_dtype: np.dtype | torch.dtype
+
+    @abstractmethod
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        pass
+
+    @abstractmethod
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+
+    @abstractmethod
+    def to(self, device: str):
+        """Move model to device."""
+
+    @abstractmethod
+    def eval(self):
+        """Set model to evaluation mode."""
+
+    @abstractmethod
+    def half(self):
+        """Convert model to half precision."""
+
+    @abstractmethod
+    def cpu(self):
+        """Move model to cpu."""
+
+    @property
+    def training(self) -> bool:
+        """Return whether model is in training mode."""
+        return False
+
+    @property
+    def device(self) -> str:
+        """Return the device of the model."""
+        return self._device
+
+    @device.setter
+    def device(self, device: str):
+        """Set the device of the model."""
+        if device == "cuda" and ":" not in device:
+            device = f"{device}:0"
+
+        self._device = device
+
+
+class TorchscriptEvaluationModel(BaseEvaluationModel):
+    """Wrapper for torchscript models."""
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.model(*args, **kwargs)
+
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+        self.model_path = model_path
+        self.device = device
+
+        model = cast(RecursiveScriptModule, torch.jit.load(self.model_path))
+        model.eval()
+        model.to(self.device)
+
+        parameter_types = {param.dtype for param in model.parameters()}
+        if len(parameter_types) == 2:
+            # TODO: There could be models with mixed precision?
+            raise ValueError(f"Expected only one type of parameters, found {parameter_types}")
+
+        self.model_dtype = list(parameter_types)[0]
+        self.model = model
+        self.is_loaded = True
+
+    def to(self, device: str):
+        """Move model to device."""
+        self.model.to(device)
+        self.device = device
+
+    def eval(self):
+        """Set model to evaluation mode."""
+        self.model.eval()
+
+    @property
+    def training(self) -> bool:
+        """Return whether model is in training mode."""
+        return self.model.training
+
+    def half(self):
+        """Convert model to half precision."""
+        self.model.half()
+
+    def cpu(self):
+        """Move model to cpu."""
+        self.model.cpu()
+
+
+class TorchEvaluationModel(TorchscriptEvaluationModel):
+    """Wrapper for torch models.
+
+    Args:
+        model_architecture: Optional torch model architecture
+    """
+
+    def __init__(self, config: DictConfig, model_architecture: nn.Module) -> None:
+        super().__init__(config=config)
+        self.model = model_architecture
+        self.model.eval()
+        device = next(self.model.parameters()).device
+        self.device = str(device)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.model(*args, **kwargs)
+
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+        self.model_path = model_path
+        self.device = device
+        self.model.load_state_dict(torch.load(self.model_path))
+        self.model.eval()
+        self.model.to(self.device)
+
+        parameter_types = {param.dtype for param in self.model.parameters()}
+        if len(parameter_types) == 2:
+            # TODO: There could be models with mixed precision?
+            raise ValueError(f"Expected only one type of parameters, found {parameter_types}")
+
+        self.model_dtype = list(parameter_types)[0]
+        self.is_loaded = True
+
+
+onnx_to_torch_dtype_dict = {
+    "tensor(bool)": torch.bool,
+    "tensor(uint8)": torch.uint8,
+    "tensor(int8)": torch.int8,
+    "tensor(int16)": torch.int16,
+    "tensor(int32)": torch.int32,
+    "tensor(int64)": torch.int64,
+    "tensor(float16)": torch.float16,
+    "tensor(float32)": torch.float32,
+    "tensor(float)": torch.float32,
+    "tensor(float64)": torch.float64,
+    "tensor(complex64)": torch.complex64,
+    "tensor(complex128)": torch.complex128,
+}
+
+
+class ONNXEvaluationModel(BaseEvaluationModel):
+    """Wrapper for ONNX models. It's designed to provide a similar interface to standard torch models."""
+
+    def __init__(self, config: DictConfig) -> None:
+        if not ONNX_AVAILABLE:
+            raise ImportError(
+                "onnxruntime is not installed. Please install ONNX capabilities for quadra with: poetry install -E onnx"
+            )
+        super().__init__(config=config)
+        self.session_options = self.generate_session_options()
+
+    def generate_session_options(self) -> ort.SessionOptions:
+        """Generate session options from the current config."""
+        session_options = ort.SessionOptions()
+
+        if hasattr(self.config, "session_options") and self.config.session_options is not None:
+            session_options_dict = cast(
+                dict[str, Any], OmegaConf.to_container(self.config.session_options, resolve=True)
+            )
+            for key, value in session_options_dict.items():
+                final_value = value
+                if isinstance(value, dict) and "_target_" in value:
+                    final_value = instantiate(final_value)
+
+                setattr(session_options, key, final_value)
+
+        return session_options
+
+    def __call__(self, *inputs: np.ndarray | torch.Tensor) -> Any:
+        """Run inference on the model and return the output as torch tensors."""
+        # TODO: Maybe we can support also kwargs
+        use_pytorch = False
+
+        onnx_inputs: dict[str, np.ndarray | torch.Tensor] = {}
+
+        for onnx_input, current_input in zip(self.model.get_inputs(), inputs):
+            if isinstance(current_input, torch.Tensor):
+                onnx_inputs[onnx_input.name] = current_input
+                use_pytorch = True
+            elif isinstance(current_input, np.ndarray):
+                onnx_inputs[onnx_input.name] = current_input
+            else:
+                raise ValueError(f"Invalid input type: {type(inputs)}")
+
+            if use_pytorch and isinstance(current_input, np.ndarray):
+                raise ValueError("Cannot mix torch and numpy inputs")
+
+        if use_pytorch:
+            onnx_output = self._forward_from_pytorch(cast(dict[str, torch.Tensor], onnx_inputs))
+        else:
+            onnx_output = self._forward_from_numpy(cast(dict[str, np.ndarray], onnx_inputs))
+
+        onnx_output = [torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x for x in onnx_output]
+
+        if len(onnx_output) == 1:
+            onnx_output = onnx_output[0]
+
+        return onnx_output
+
+    def _forward_from_pytorch(self, input_dict: dict[str, torch.Tensor]):
+        """Run inference on the model and return the output as torch tensors."""
+        io_binding = self.model.io_binding()
+        device_type = self.device.split(":")[0]
+
+        for k, v in input_dict.items():
+            if not v.is_contiguous():
+                # If not contiguous onnx give wrong results
+                v = v.contiguous()  # noqa: PLW2901
+
+            io_binding.bind_input(
+                name=k,
+                device_type=device_type,
+                # Weirdly enough onnx wants 0 for cpu
+                device_id=0 if device_type == "cpu" else int(self.device.split(":")[1]),
+                element_type=np.float16 if v.dtype == torch.float16 else np.float32,
+                shape=tuple(v.shape),
+                buffer_ptr=v.data_ptr(),
+            )
+
+        for x in self.model.get_outputs():
+            # TODO: Is it possible to also bind the output? We require info about output dimensions
+            io_binding.bind_output(name=x.name)
+
+        self.model.run_with_iobinding(io_binding)
+
+        output = io_binding.copy_outputs_to_cpu()
+
+        return output
+
+    def _forward_from_numpy(self, input_dict: dict[str, np.ndarray]):
+        """Run inference on the model and return the output as numpy array."""
+        ort_outputs = [x.name for x in self.model.get_outputs()]
+
+        onnx_output = self.model.run(ort_outputs, input_dict)
+
+        return onnx_output
+
+    def load_from_disk(self, model_path: str, device: str = "cpu"):
+        """Load model from disk."""
+        self.model_path = model_path
+        self.device = device
+
+        ort_providers = self._get_providers(device)
+        self.model = ort.InferenceSession(self.model_path, providers=ort_providers, sess_options=self.session_options)
+        self.model_dtype = self.cast_onnx_dtype(self.model.get_inputs()[0].type)
+        self.is_loaded = True
+
+    def _get_providers(self, device: str) -> list[tuple[str, dict[str, Any]] | str]:
+        """Return the providers for the ONNX model based on the device."""
+        ort_providers: list[tuple[str, dict[str, Any]] | str]
+
+        if device == "cpu":
+            ort_providers = ["CPUExecutionProvider"]
+        else:
+            ort_providers = [
+                (
+                    "CUDAExecutionProvider",
+                    {
+                        "device_id": int(device.split(":")[1]),
+                    },
+                )
+            ]
+
+        return ort_providers
+
+    def to(self, device: str):
+        """Move model to device."""
+        self.device = device
+        ort_providers = self._get_providers(device)
+        self.model.set_providers(ort_providers)
+
+    def eval(self):
+        """Fake interface to match torch models."""
+        return self
+
+    def half(self):
+        """Convert model to half precision."""
+        raise NotImplementedError("At the moment ONNX models do not support half method.")
+
+    def cpu(self):
+        """Move model to cpu."""
+        self.to("cpu")
+
+    def cast_onnx_dtype(self, onnx_dtype: str) -> torch.dtype | np.dtype:
+        """Cast ONNX dtype to numpy or pytorch dtype."""
+        return onnx_to_torch_dtype_dict[onnx_dtype]

quadra/modules/backbone.py

@@ -0,0 +1,30 @@
+from typing import Any
+
+import segmentation_models_pytorch as smp
+
+
+def create_smp_backbone(
+    arch: str,
+    encoder_name: str,
+    freeze_encoder: bool = False,
+    in_channels: int = 3,
+    num_classes: int = 0,
+    **kwargs: Any,
+):
+    """Create Segmentation.models.pytorch model backbone
+    Args:
+        arch: architecture name
+        encoder_name: architecture name
+        freeze_encoder: freeze encoder or not
+        in_channels: number of input channels
+        num_classes: number of classes
+        **kwargs: extra arguments for model (for example classification head).
+    """
+    model = smp.create_model(
+        arch=arch, encoder_name=encoder_name, in_channels=in_channels, classes=num_classes, **kwargs
+    )
+    if freeze_encoder:
+        for child in model.encoder.children():
+            for param in child.parameters():
+                param.requires_grad = False
+    return model

quadra/modules/base.py

@@ -0,0 +1,312 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import pytorch_lightning as pl
+import sklearn
+import torch
+import torchmetrics
+from sklearn.linear_model import LogisticRegression
+from torch import nn
+from torch.optim import Optimizer
+
+from quadra.models.base import ModelSignatureWrapper
+
+__all__ = ["BaseLightningModule", "SSLModule"]
+
+
+class BaseLightningModule(pl.LightningModule):
+    """Base lightning module.
+
+    Args:
+        model: Network Module used for extract features
+        optimizer: optimizer of the training. If None a default Adam is used.
+        lr_scheduler: lr scheduler. If None a default ReduceLROnPlateau is used.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        optimizer: Optimizer | None = None,
+        lr_scheduler: object | None = None,
+        lr_scheduler_interval: str | None = "epoch",
+    ):
+        super().__init__()
+        self.model = ModelSignatureWrapper(model)
+        self.optimizer = optimizer
+        self.schedulers = lr_scheduler
+        self.lr_scheduler_interval = lr_scheduler_interval
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward method
+        Args:
+            x: input tensor.
+
+        Returns:
+            model inference
+        """
+        return self.model(x)
+
+    def configure_optimizers(self) -> tuple[list[Any], list[dict[str, Any]]]:
+        """Get default optimizer if not passed a value.
+
+        Returns:
+            optimizer and lr scheduler as Tuple containing a list of optimizers and a list of lr schedulers
+        """
+        # get default optimizer
+        if getattr(self, "optimizer", None) is None or not self.optimizer:
+            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=3e-4)
+
+        # get default scheduler
+        if getattr(self, "schedulers", None) is None or not self.schedulers:
+            self.schedulers = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, factor=0.5, patience=30)
+
+        lr_scheduler_conf = {
+            "scheduler": self.schedulers,
+            "interval": self.lr_scheduler_interval,
+            "monitor": "val_loss",
+            "strict": False,
+        }
+        return [self.optimizer], [lr_scheduler_conf]
+
+    # pylint: disable=unused-argument
+    def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx: int = 0):
+        """Redefine optimizer zero grad."""
+        optimizer.zero_grad(set_to_none=True)
+
+
+class SSLModule(BaseLightningModule):
+    """Base module for self supervised learning.
+
+    Args:
+        model: Network Module used for extract features
+        criterion: SSL loss to be applied
+        classifier: Standard sklearn classifiers
+        optimizer: optimizer of the training. If None a default Adam is used.
+        lr_scheduler: lr scheduler. If None a default ReduceLROnPlateau is used.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        criterion: nn.Module,
+        classifier: sklearn.base.ClassifierMixin | None = None,
+        optimizer: Optimizer | None = None,
+        lr_scheduler: object | None = None,
+        lr_scheduler_interval: str | None = "epoch",
+    ):
+        super().__init__(model, optimizer, lr_scheduler, lr_scheduler_interval)
+        self.criterion = criterion
+        self.classifier_train_loader: torch.utils.data.DataLoader | None
+        if classifier is None:
+            self.classifier = LogisticRegression(max_iter=10000, n_jobs=8, random_state=42)
+        else:
+            self.classifier = classifier
+
+        self.val_acc = torchmetrics.Accuracy()
+
+    def fit_estimator(self):
+        """Fit a classifier on the embeddings extracted from the current trained model."""
+        targets = []
+        train_embeddings = []
+        self.model.eval()
+        with torch.no_grad():
+            for im, target in self.classifier_train_loader:
+                emb = self.model(im.to(self.device))
+                targets.append(target)
+                train_embeddings.append(emb)
+        targets = torch.cat(targets, dim=0).cpu().numpy()
+        train_embeddings = torch.cat(train_embeddings, dim=0).cpu().numpy()
+        self.classifier.fit(train_embeddings, targets)
+
+    def calculate_accuracy(self, batch):
+        """Calculate accuracy on a batch of data."""
+        images, labels = batch
+        with torch.no_grad():
+            embedding = self.model(images).cpu().numpy()
+
+        predictions = self.classifier.predict(embedding)
+        labels = labels.detach()
+        acc = self.val_acc(torch.tensor(predictions, device=self.device), labels)
+
+        return acc
+
+    # TODO: In multiprocessing mode, this function is called multiple times, how can we avoid this?
+    def on_validation_start(self) -> None:
+        if not hasattr(self, "classifier_train_loader") and hasattr(self.trainer, "datamodule"):
+            self.classifier_train_loader = self.trainer.datamodule.classifier_train_dataloader()
+
+        if self.classifier_train_loader is not None:
+            self.fit_estimator()
+
+    def validation_step(self, batch: tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor], batch_idx: int) -> None:
+        # pylint: disable=unused-argument
+        if self.classifier_train_loader is None:
+            # Compute loss
+            (im_x, im_y), _ = batch
+            z1 = self(im_x)
+            z2 = self(im_y)
+            loss = self.criterion(z1, z2)
+
+            self.log(
+                "val_loss",
+                loss,
+                on_epoch=True,
+                on_step=True,
+                logger=True,
+                prog_bar=True,
+            )
+            return loss
+
+        acc = self.calculate_accuracy(batch)
+        self.log("val_acc", acc, on_epoch=True, on_step=False, logger=True, prog_bar=True)
+        return None
+
+
+class SegmentationModel(BaseLightningModule):
+    """Generic segmentation model.
+
+    Args:
+        model: segmentation model to be used.
+        loss_fun: loss function to be used.
+        optimizer: Optimizer to be used. Defaults to None.
+        lr_scheduler: lr scheduler to be used. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        loss_fun: Callable,
+        optimizer: Optimizer | None = None,
+        lr_scheduler: object | None = None,
+    ):
+        super().__init__(model, optimizer, lr_scheduler)
+        self.loss_fun = loss_fun
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward method
+        Args:
+            x: input tensor.
+
+        Returns:
+            model inference
+        """
+        x = self.model(x)
+        return x
+
+    def step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute loss
+        Args:
+            batch: batch.
+
+        Returns:
+            Prediction and target masks
+        """
+        images, target_masks, _ = batch
+        pred_masks = self(images)
+        if len(pred_masks.shape) == 3:
+            pred_masks = pred_masks.unsqueeze(1)
+        if len(target_masks.shape) == 3:
+            target_masks = target_masks.unsqueeze(1)
+        assert pred_masks.shape == target_masks.shape
+
+        return pred_masks, target_masks
+
+    def compute_loss(self, pred_masks: torch.Tensor, target_masks: torch.Tensor) -> torch.Tensor:
+        """Compute loss
+        Args:
+            pred_masks: predicted masks
+            target_masks: target masks.
+
+        Returns:
+            The computed loss
+
+        """
+        loss = self.loss_fun(pred_masks, target_masks)
+        return loss
+
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int):
+        """Training step."""
+        # pylint: disable=unused-argument
+        pred_masks, target_masks = self.step(batch)
+        loss = self.compute_loss(pred_masks, target_masks)
+        self.log_dict(
+            {"loss": loss},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx):
+        """Validation step."""
+        # pylint: disable=unused-argument
+        pred_masks, target_masks = self.step(batch)
+        loss = self.compute_loss(pred_masks, target_masks)
+        self.log_dict(
+            {"val_loss": loss},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def test_step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor], batch_idx: int):
+        """Test step."""
+        # pylint: disable=unused-argument
+        pred_masks, target_masks = self.step(batch)
+        loss = self.compute_loss(pred_masks, target_masks)
+        self.log_dict(
+            {"test_loss": loss},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def predict_step(
+        self,
+        batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+        batch_idx: int,
+        dataloader_idx: int | None = None,
+    ) -> Any:
+        """Predict step."""
+        # pylint: disable=unused-argument
+        images, masks, labels = batch
+        pred_masks = self(images)
+        return images.cpu(), masks.cpu(), pred_masks.cpu(), labels.cpu()
+
+
+class SegmentationModelMulticlass(SegmentationModel):
+    """Generic multiclass segmentation model.
+
+    Args:
+        model: segmentation model to be used.
+        loss_fun: loss function to be used.
+        optimizer: Optimizer to be used. Defaults to None.
+        lr_scheduler: lr scheduler to be used. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        loss_fun: Callable,
+        optimizer: Optimizer | None = None,
+        lr_scheduler: object | None = None,
+    ):
+        super().__init__(model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, loss_fun=loss_fun)
+
+    def step(self, batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute step
+        Args:
+            batch: batch.
+
+        Returns:
+            prediction, target
+
+        """
+        images, target_masks, _ = batch
+        pred_masks = self(images)
+
+        return pred_masks, target_masks

quadra/modules/classification/__init__.py

@@ -0,0 +1,3 @@
+from .base import ClassificationModule, MultilabelClassificationModule
+
+__all__ = ["ClassificationModule", "MultilabelClassificationModule"]