From 4e9354a46cc71615be8b7896ebdeb802f87836f2 Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 5 Dec 2025 23:00:51 +0100 Subject: [PATCH 01/10] Adding VAEABMIL --- docs/api/index.md | 3 + docs/api/models/index.md | 1 + docs/api/models/vaeabmil.md | 10 + docs/api/nn/index.md | 2 + docs/api/nn/mlp.md | 6 + docs/api/nn/variational_autoencoder.md | 24 + mkdocs.yml | 3 + tests/models/test_vaeabmil.py | 87 +++ tests/nn/test_mlp.py | 86 +++ tests/nn/test_variational_autoencoder.py | 143 +++++ torchmil/datasets/processed_mil_dataset.py | 1 - torchmil/models/__init__.py | 1 + torchmil/models/vaeabmil.py | 168 ++++++ torchmil/nn/__init__.py | 4 + torchmil/nn/mlp.py | 61 +++ torchmil/nn/variational_autoencoder.py | 604 +++++++++++++++++++++ 16 files changed, 1203 insertions(+), 1 deletion(-) create mode 100644 docs/api/models/vaeabmil.md create mode 100644 docs/api/nn/mlp.md create mode 100644 docs/api/nn/variational_autoencoder.md create mode 100644 tests/models/test_vaeabmil.py create mode 100644 tests/nn/test_mlp.py create mode 100644 tests/nn/test_variational_autoencoder.py create mode 100644 torchmil/models/vaeabmil.py create mode 100644 torchmil/nn/mlp.py create mode 100644 torchmil/nn/variational_autoencoder.py diff --git a/docs/api/index.md b/docs/api/index.md index c85642e..ddbe4fb 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -45,6 +45,8 @@ - [Sm operator](nn/sm.md) - [Max Pool](nn/max_pool.md) - [Mean Pool](nn/mean_pool.md) +- [Multi Layer Perceptron](nn/mlp.md) +- [Variational Autoencoder](nn/variational_autoencoder.md) ## Models: [torchmil.models](models/index.md) - [Introduction](models/index.md) @@ -65,6 +67,7 @@ - [TransformerProbSmoothABMIL](models/transformer_prob_smooth_abmil.md) - [TransMIL](models/transmil.md) - [SETMIL](models/setmil.md) +- [VAEABMIL](models/vaeabmil.md) ## Visualize: [torchmil.visualize](visualize/index.md) - [Introduction](visualize/index.md) diff --git a/docs/api/models/index.md b/docs/api/models/index.md index 508bbf3..38e19ae 100644 --- a/docs/api/models/index.md +++ b/docs/api/models/index.md @@ -19,3 +19,4 @@ They all inherit from the [General MIL model](mil_model.md) class, which provide - [TransMIL](transmil.md) - [SETMIL](setmil.md) - [IIBMIL](iibmil.md) +- [VAEABMIL](vaeabmil.md) diff --git a/docs/api/models/vaeabmil.md b/docs/api/models/vaeabmil.md new file mode 100644 index 0000000..f6d1c09 --- /dev/null +++ b/docs/api/models/vaeabmil.md @@ -0,0 +1,10 @@ +# VAEABMIL + +::: torchmil.models.VAEABMIL + options: + members: + - __init__ + - forward + - compute_loss + - predict + - log_marginal_likelihood_importance_sampling diff --git a/docs/api/nn/index.md b/docs/api/nn/index.md index 1fc8d9e..089064c 100644 --- a/docs/api/nn/index.md +++ b/docs/api/nn/index.md @@ -27,3 +27,5 @@ These modules are designed to be flexible and easy to use, allowing you to build - [Sm operator](sm.md) - [Max Pool](max_pool.md) - [Mean Pool](mean_pool.md) +- [Multi Layer Perceptron](mlp.md) +- [Variational Autoencoder](variational_autoencoder.md) diff --git a/docs/api/nn/mlp.md b/docs/api/nn/mlp.md new file mode 100644 index 0000000..5c4a006 --- /dev/null +++ b/docs/api/nn/mlp.md @@ -0,0 +1,6 @@ +# MLP +::: torchmil.nn.MLP + options: + members: + - __init__ + - forward diff --git a/docs/api/nn/variational_autoencoder.md b/docs/api/nn/variational_autoencoder.md new file mode 100644 index 0000000..08e39af --- /dev/null +++ b/docs/api/nn/variational_autoencoder.md @@ -0,0 +1,24 @@ +# Variational Autoencoder +::: torchmil.nn.VariationalAutoEncoder + options: + members: + - __init__ + - get_reparameterized_samples + - get_raw_output_enc + - get_raw_output_dec + - forward + - get_posterior_samples + - complete_forward_samples + - compute_loss + - _kl_prior + - _diagonal_log_gaussian_pdf + - log_marginal_likelihood_importance_sampling +------------------------- +::: torchmil.nn.VariationalAutoEncoderMIL + options: + members: + - __init__ + - forward + - log_marginal_likelihood_importance_sampling + - compute_loss + - complete_forward_samples diff --git a/mkdocs.yml b/mkdocs.yml index a4d9b9d..3fb5491 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -168,6 +168,8 @@ nav: - api/nn/sm.md - api/nn/mean_pool.md - api/nn/max_pool.md + - api/nn/mlp.md + - api/nn/variational_autoencoder.md - torchmil.models: - api/models/index.md - api/models/mil_model.md @@ -188,6 +190,7 @@ nav: - api/models/transmil.md - api/models/setmil.md - api/models/iibmil.md + - api/models/vaeabmil.md - torchmil.visualize: - api/visualize/index.md - api/visualize/vis_ctscan.md diff --git a/tests/models/test_vaeabmil.py b/tests/models/test_vaeabmil.py new file mode 100644 index 0000000..cde7b63 --- /dev/null +++ b/tests/models/test_vaeabmil.py @@ -0,0 +1,87 @@ +import torch +import pytest + +from torchmil.models import VAEABMIL # Import the VAEABMIL class +from torchmil.nn import VariationalAutoEncoderMIL + + +# Fixtures for common setup +@pytest.fixture +def sample_data(): + # Returns a tuple of (X, Y, mask) + torch.manual_seed(42) # For reproducibility + X = torch.randn(2, 3, 10) # batch_size, bag_size, feat_dim + Y = torch.randint(0, 2, (2,)).float() # batch_size, ensure float for BCE loss + mask = torch.ones(2, 3).bool() # All instances are valid for more stable testing + return X, Y, mask + + +@pytest.fixture +def vae_feat_ext(): + # Returns a VariationalAutoEncoderMIL instance for feature extraction + return VariationalAutoEncoderMIL( + input_shape=(10,), + layer_sizes=[8, 5], + activations=["relu", "None"] + ) + + +@pytest.fixture +def vaeabmil_model(vae_feat_ext): + # Returns an instance of the VAEABMIL model with default parameters + return VAEABMIL( + feat_ext=vae_feat_ext, + in_shape=(3, 10) + ) + + +# Basic tests for VAEABMIL class +def test_vaeabmil_initialization(vae_feat_ext): + # Test basic initialization + model = VAEABMIL(feat_ext=vae_feat_ext, in_shape=(3, 10)) + assert model is not None + + +def test_vaeabmil_forward_pass(sample_data, vaeabmil_model): + # Test basic forward pass + X, _, mask = sample_data + + Y_pred = vaeabmil_model(X, mask) + assert Y_pred.shape == (2,) + + +def test_vaeabmil_forward_with_attention(sample_data, vaeabmil_model): + # Test forward pass with attention return + X, _, mask = sample_data + + Y_pred, att = vaeabmil_model(X, mask, return_att=True) + assert Y_pred.shape == (2,) + assert att.shape == (2, 3) + + +def test_vaeabmil_compute_loss(sample_data, vaeabmil_model): + # Test loss computation + X, Y, mask = sample_data + + # Ensure the model is in a good state for loss computation + with torch.no_grad(): + _ = vaeabmil_model(X, mask) + + Y_pred, loss_dict = vaeabmil_model.compute_loss(Y, X, mask) + + assert Y_pred.shape == (2,) + assert "BCEWithLogitsLoss" in loss_dict + assert "VaeELL" in loss_dict + assert "VaeKL" in loss_dict + + +def test_vaeabmil_predict(sample_data, vaeabmil_model): + # Test predict method + X, _, mask = sample_data + + Y_pred = vaeabmil_model.predict(X, mask, return_inst_pred=False) + assert Y_pred.shape == (2,) + + Y_pred, y_inst_pred = vaeabmil_model.predict(X, mask, return_inst_pred=True) + assert Y_pred.shape == (2,) + assert y_inst_pred.shape == (2, 3) \ No newline at end of file diff --git a/tests/nn/test_mlp.py b/tests/nn/test_mlp.py new file mode 100644 index 0000000..a8d46cc --- /dev/null +++ b/tests/nn/test_mlp.py @@ -0,0 +1,86 @@ +import torch +import pytest + +from torchmil.nn.mlp import MLP, get_activation + + +# Test activation function utility +def test_get_activation(): + # Test different activation functions + relu = get_activation("relu") + assert isinstance(relu, torch.nn.ReLU) + + sigmoid = get_activation("sigmoid") + assert isinstance(sigmoid, torch.nn.Sigmoid) + + tanh = get_activation("tanh") + assert isinstance(tanh, torch.nn.Tanh) + + # Test None case + none_activation = get_activation("none") + assert none_activation is None + + +# Fixtures for common setup +@pytest.fixture +def sample_input(): + return torch.randn(5, 10) # batch_size=5, input_dim=10 + + +@pytest.fixture +def mlp_basic(): + return MLP(input_size=10, linear_sizes=[5, 3], activations=["relu", "relu"]) + + +# Basic tests for MLP class +def test_mlp_initialization(): + # Test basic initialization + mlp = MLP(input_size=10, linear_sizes=[5, 3], activations=["relu", "relu"]) + assert mlp.input_size == 10 + assert mlp.linear_sizes == [5, 3] + assert mlp.activations == ["relu", "relu"] + + +def test_mlp_initialization_default_params(): + # Test initialization with default parameters + mlp = MLP() + assert mlp.input_size == 512 + assert mlp.linear_sizes == [100, 50] + assert mlp.activations == ["relu", "relu"] + + +def test_mlp_forward_pass(sample_input, mlp_basic): + # Test basic forward pass + output = mlp_basic(sample_input) + assert output.shape == (5, 3) # batch_size, final_layer_size + + +def test_mlp_no_activation(): + # Test with no activation (None activation) + mlp = MLP(input_size=10, linear_sizes=[5], activations=["None"]) + input_data = torch.randn(2, 10) + output = mlp(input_data) + assert output.shape == (2, 5) + + +def test_mlp_different_activations(): + # Test different activation functions work + input_data = torch.randn(3, 10) + + # ReLU activation + mlp_relu = MLP(input_size=10, linear_sizes=[5], activations=["relu"]) + output_relu = mlp_relu(input_data) + assert output_relu.shape == (3, 5) + + # Sigmoid activation + mlp_sigmoid = MLP(input_size=10, linear_sizes=[5], activations=["sigmoid"]) + output_sigmoid = mlp_sigmoid(input_data) + assert output_sigmoid.shape == (3, 5) + + +def test_mlp_multiple_layers(): + # Test MLP with multiple layers + mlp = MLP(input_size=20, linear_sizes=[16, 8, 4], activations=["relu", "relu", "None"]) + input_data = torch.randn(2, 20) + output = mlp(input_data) + assert output.shape == (2, 4) \ No newline at end of file diff --git a/tests/nn/test_variational_autoencoder.py b/tests/nn/test_variational_autoencoder.py new file mode 100644 index 0000000..32e6fb4 --- /dev/null +++ b/tests/nn/test_variational_autoencoder.py @@ -0,0 +1,143 @@ +import torch +import pytest + +from torchmil.nn.variational_autoencoder import VariationalAutoEncoder, VariationalAutoEncoderMIL + + +# Fixtures for common setup +@pytest.fixture +def sample_data(): + return torch.randn(4, 10) # batch_size=4, input_dim=10 + + +@pytest.fixture +def sample_bag_data(): + return torch.randn(2, 3, 10) # batch_size=2, bag_size=3, input_dim=10 + + +@pytest.fixture +def vae_basic(): + return VariationalAutoEncoder( + input_shape=(10,), + layer_sizes=[8, 5], + activations=["relu", "None"] + ) + + +@pytest.fixture +def vae_mil(): + return VariationalAutoEncoderMIL( + input_shape=(10,), + layer_sizes=[8, 5], + activations=["relu", "None"] + ) + + +# Basic tests for VariationalAutoEncoder class +def test_vae_initialization(): + # Test basic initialization + vae = VariationalAutoEncoder( + input_shape=(10,), + layer_sizes=[8, 5], + activations=["relu", "None"] + ) + assert vae.input_dim == (10,) + assert vae.output_size == 10 + assert vae.layer_sizes == [8, 5] + + +def test_vae_initialization_diagonal_covar(): + # Test initialization with diagonal covariance + vae = VariationalAutoEncoder( + input_shape=(10,), + layer_sizes=[8, 5], + covar_mode="diagonal" + ) + assert vae.covar_mode == "diagonal" + + +def test_vae_initialization_invalid_covar(): + # Test that invalid covariance mode raises error + with pytest.raises(NotImplementedError): + VariationalAutoEncoder( + input_shape=(10,), + layer_sizes=[8, 5], + covar_mode="invalid" + ) + + +def test_vae_forward(sample_data, vae_basic): + # Test forward pass (encoding only) + samples = vae_basic(sample_data, n_samples=2) + assert samples.shape == (4, 2, 5) # batch_size, n_samples, latent_dim + + +def test_vae_get_posterior_samples(sample_data, vae_basic): + # Test posterior sampling + samples = vae_basic.get_posterior_samples(sample_data, n_samples=2) + assert samples.shape == (4, 2, 5) # batch_size, n_samples, latent_dim + + +def test_vae_complete_forward_samples(sample_data, vae_basic): + # Test complete forward pass (encode + decode) + reconstructions = vae_basic.complete_forward_samples(sample_data, n_samples=1) + assert reconstructions.shape == sample_data.shape + + +def test_vae_compute_loss(sample_data, vae_basic): + # Test loss computation + loss_dict = vae_basic.compute_loss(sample_data, reduction="sum", n_samples=2) + + assert "VaeELL" in loss_dict + assert "VaeKL" in loss_dict + assert loss_dict["VaeELL"].shape == () # scalar + assert loss_dict["VaeKL"].shape == () # scalar + + +def test_vae_get_raw_output_enc(sample_data, vae_basic): + # Test encoder raw output + mean, log_std = vae_basic.get_raw_output_enc(sample_data) + + assert mean.shape == (4, 5) # batch_size, latent_dim + assert log_std.shape == (4, 1) # batch_size, d_var_enc (single mode) + + +def test_vae_get_raw_output_dec(vae_basic): + # Test decoder raw output + latent_samples = torch.randn(4, 5) # batch_size, latent_dim + mean, log_std = vae_basic.get_raw_output_dec(latent_samples) + + assert mean.shape == (4, 10) # batch_size, input_dim + # In single covar mode, log_std is expanded to match input dim + assert log_std.shape == (4, 10) # batch_size, input_dim (expanded from d_var_dec) + + +# Basic tests for VariationalAutoEncoderMIL class +def test_vae_mil_initialization(): + # Test MIL VAE initialization + vae_mil = VariationalAutoEncoderMIL( + input_shape=(10,), + layer_sizes=[8, 5] + ) + assert isinstance(vae_mil, VariationalAutoEncoder) + + +def test_vae_mil_forward(sample_bag_data, vae_mil): + # Test MIL VAE forward pass + samples = vae_mil(sample_bag_data, n_samples=2) + assert samples.shape == (2, 3, 2, 5) # batch_size, bag_size, n_samples, latent_dim + + +def test_vae_mil_compute_loss(sample_bag_data, vae_mil): + # Test MIL VAE loss computation + loss_dict = vae_mil.compute_loss(sample_bag_data, reduction="mean") + + assert "VaeELL" in loss_dict and "VaeKL" in loss_dict + assert loss_dict["VaeELL"].shape == () + assert loss_dict["VaeKL"].shape == () + + +def test_vae_mil_complete_forward_samples(sample_bag_data, vae_mil): + # Test complete forward pass for MIL VAE + reconstructions = vae_mil.complete_forward_samples(sample_bag_data) + assert reconstructions.shape == sample_bag_data.shape \ No newline at end of file diff --git a/torchmil/datasets/processed_mil_dataset.py b/torchmil/datasets/processed_mil_dataset.py index f4a0e7d..14dcd9e 100644 --- a/torchmil/datasets/processed_mil_dataset.py +++ b/torchmil/datasets/processed_mil_dataset.py @@ -157,7 +157,6 @@ def __init__( self.loaded_bags = {} if self.load_at_init: for name in self.bag_names: - print("name") self.loaded_bags[name] = self._build_bag(name) def _set_file_read_fn(self, file_type: str) -> None: diff --git a/torchmil/models/__init__.py b/torchmil/models/__init__.py index c99e8f3..84a0e9f 100644 --- a/torchmil/models/__init__.py +++ b/torchmil/models/__init__.py @@ -26,3 +26,4 @@ from .iibmil import IIBMIL as IIBMIL from .setmil import SETMIL as SETMIL from .gtp import GTP as GTP +from .vaeabmil import VAEABMIL as VAEABMIL \ No newline at end of file diff --git a/torchmil/models/vaeabmil.py b/torchmil/models/vaeabmil.py new file mode 100644 index 0000000..6702299 --- /dev/null +++ b/torchmil/models/vaeabmil.py @@ -0,0 +1,168 @@ +import torch + +from .mil_model import MILModel +from torchmil.nn import AttentionPool, LazyLinear +from torchmil.nn.utils import get_feat_dim +from torchmil.nn import VariationalAutoEncoderMIL + + +class VAEABMIL(MILModel): + r""" + Variational Autoencoder - Attention-based Multiple Instance Learning (VAEABMIL) model, proposed in the paper [Using Variational Autoencoders for Out of Distribution Detection in Histological Multiple Instance Learning](https://ieeexplore.ieee.org/abstract/document/11098836/). + + The model jointly trains a Variational Autoencoder (VAE) on instance features to learn a latent representation that is used for attention-based multiple instance learning and to detect out-of-distribution instances and bags. + + Given an input bag $\mathbf{X} = \left[ \mathbf{x}_1, \ldots, \mathbf{x}_N \right]^\top \in \mathbb{R}^{N \times P}$, the model uses the VAE, to obtain an approximated posterior distribution $p(\mathbf{z} | \mathbf{x})$ for each instance $\mathbf{x}$ in the bag. Then, $\mathbf{X} = [\mathbf{z}_1, \ldots, \mathbf{z}_N] \in \mathbb{R}^{N \times D}$ with $\mathbf{z}_i \sim p(\mathbf{z}_i \mid \mathbf{x}_i)$. + + Lastly, it aggregates the instance features into a bag representation $\mathbf{z} \in \mathbb{R}^{D}$ using the attention-based pooling, + + $$ + \mathbf{z}, \mathbf{f} = \operatorname{AttentionPool}(\mathbf{X}). + $$ + + where $\mathbf{f} \in \mathbb{R}^{N}$ are the attention values. + See [AttentionPool](../nn/attention/attention_pool.md) for more details on the attention-based pooling. + The bag representation $\mathbf{z}$ is then fed into a classifier (one linear layer) to predict the bag label. + """ + + def __init__( + self, + feat_ext: VariationalAutoEncoderMIL, + in_shape: tuple = None, + att_dim: int = 128, + att_act: str = "tanh", + gated: bool = False, + criterion: torch.nn.Module = torch.nn.BCEWithLogitsLoss(), + vae_loss_reduction: str = "mean", + ) -> None: + """ + Arguments: + feat_ext: Variational Autoencoder used as feature extractor. + in_shape: Shape of input data expected by the feature extractor (excluding batch dimension). If not provided, it will be lazily initialized. + att_dim: Attention dimension. + att_act: Activation function for attention. Possible values: 'tanh', 'relu', 'gelu'. + gated: If True, use gated attention in the attention pooling. + criterion: Loss function. By default, Binary Cross-Entropy loss from logits. + vae_loss_reduction: Reduction method for VAE loss. Possible values: 'sum', 'mean', 'none'. + """ + super().__init__() + self.criterion = criterion + self.vae_loss_reduction = vae_loss_reduction + + self.feat_ext = feat_ext + if in_shape is not None: + feat_dim = get_feat_dim(feat_ext, in_shape) + else: + feat_dim = None + + self.pool = AttentionPool( + in_dim=feat_dim, att_dim=att_dim, act=att_act, gated=gated + ) + + self.classifier = LazyLinear(in_features=feat_dim, out_features=1) + + def forward( + self, + X: torch.Tensor, + mask: torch.Tensor = None, + return_att: bool = False, + return_latent_repr: bool = False, + n_samples: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass. + + Arguments: + X: Bag features of shape `(batch_size, bag_size, ...)`. + mask: Mask of shape `(batch_size, bag_size)`. + return_att: If True, returns attention values (before normalization) in addition to `Y_pred`. + return_latent_repr: If True, returns latent representation in addition to `Y_pred`. (Currently not implemented) + n_samples: Number of Monte Carlo samples to use for the VAE. + + Returns: + Y_pred: Bag label logits of shape `(batch_size,)`. + att: Only returned when `return_att=True`. Attention values (before normalization) of shape `(batch_size, bag_size)`. + """ + + X = self.feat_ext(X, n_samples) # (batch_size, bag_size, n_samples, feat_dim) + X = X.mean(dim=2) # (batch_size, bag_size, feat_dim) + + out_pool = self.pool(X, mask, return_att) # (batch_size, feat_dim) + + if return_att: + Z, f = out_pool # (batch_size, feat_dim), (batch_size, bag_size) + else: + Z = out_pool # (batch_size, feat_dim) + + Y_pred = self.classifier(Z) # (batch_size, 1) + Y_pred = Y_pred.squeeze(-1) # (batch_size,) + + if return_att: + return Y_pred, f + else: + return Y_pred + + def compute_loss( + self, + Y: torch.Tensor, + X: torch.Tensor, + mask: torch.Tensor = None, + n_samples: int = 1, + ) -> tuple[torch.Tensor, dict]: + """ + Compute loss given true bag labels. + + Arguments: + Y: Bag labels of shape `(batch_size,)`. + X: Bag features of shape `(batch_size, bag_size, ...)`. + mask: Mask of shape `(batch_size, bag_size)`. + n_samples: Number of Monte Carlo samples to use for the VAE loss computation. + + Returns: + Y_pred: Bag label logits of shape `(batch_size,)`. + loss_dict: Dictionary containing the loss values. Includes the main criterion loss and VAE losses (VaeELL and VaeKL). + """ + + Y_pred = self.forward(X, mask, return_att=False) + vae_loss = self.feat_ext.compute_loss( + X, n_samples=n_samples, reduction=self.vae_loss_reduction + ) + + crit_loss = self.criterion(Y_pred.float(), Y.float()) + crit_name = self.criterion.__class__.__name__ + + return Y_pred, {crit_name: crit_loss, **vae_loss} + + def predict( + self, X: torch.Tensor, mask: torch.Tensor = None, return_inst_pred: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Predict bag and (optionally) instance labels. + + Arguments: + X: Bag features of shape `(batch_size, bag_size, ...)`. + mask: Mask of shape `(batch_size, bag_size)`. + return_inst_pred: If `True`, returns instance labels predictions (attention values), in addition to bag label predictions. + + Returns: + Y_pred: Bag label logits of shape `(batch_size,)`. + y_inst_pred: If `return_inst_pred=True`, returns instance labels predictions (attention values) of shape `(batch_size, bag_size)`. + """ + return self.forward(X, mask, return_att=return_inst_pred) + + def log_marginal_likelihood_importance_sampling( + self, X: torch.Tensor, mask: torch.Tensor = None, n_samples: int = 1 + ) -> torch.Tensor: + """ + Estimate the marginal log-likelihood of the input bag using importance sampling. + + Arguments: + X: Bag features of shape `(batch_size, bag_size, ...)`. + n_samples: Number of importance samples to use. + + Returns: + log_likelihood: Estimated marginal log-likelihood of shape `(batch_size,)`. + """ + return self.feat_ext.log_marginal_likelihood_importance_sampling( + X, mask=mask, n_samples=n_samples + ) diff --git a/torchmil/nn/__init__.py b/torchmil/nn/__init__.py index 764bdaf..1a377df 100644 --- a/torchmil/nn/__init__.py +++ b/torchmil/nn/__init__.py @@ -39,3 +39,7 @@ ChebConv as ChebConv, dense_mincut_pool as dense_mincut_pool, ) + +from .mlp import MLP as MLP +from .variational_autoencoder import (VariationalAutoEncoder as VariationalAutoEncoder, + VariationalAutoEncoderMIL as VariationalAutoEncoderMIL) \ No newline at end of file diff --git a/torchmil/nn/mlp.py b/torchmil/nn/mlp.py new file mode 100644 index 0000000..6b87d2b --- /dev/null +++ b/torchmil/nn/mlp.py @@ -0,0 +1,61 @@ +import torch + + +def get_activation(name): + """ + Get torch activation function by name. + """ + if "relu" in name: + return torch.nn.ReLU() + elif "sigmoid" in name: + return torch.nn.Sigmoid() + elif "tanh" in name: + return torch.nn.Tanh() + else: + return None + + +class MLP(torch.nn.Module): + """ + Multi-Layer Perceptron (MLP) class. + """ + + def __init__( + self, input_size=512, linear_sizes=[100, 50], activations=["relu", "relu"] + ) -> None: + """ + Arguments: + input_size (int): Size of the input features. + linear_sizes (list of int): List containing the sizes of each linear layer. + activations (list of str): List containing the activation functions for each layer. + """ + super(MLP, self).__init__() + + self.input_size = input_size + self.linear_sizes = linear_sizes + self.activations = activations + layers = [torch.nn.Linear(self.input_size, linear_sizes[0])] + if activations[0] not in ["None", "none"]: + layers.append(get_activation(activations[0])) + + for i in range(1, len(linear_sizes)): + layers.append(torch.nn.Linear(linear_sizes[i - 1], linear_sizes[i])) + if activations[i] not in ["None", "none"]: + layers.append(get_activation(activations[i])) + + # Filter out None values that might have been added + layers = [layer for layer in layers if layer is not None] + self.net = torch.nn.Sequential(*layers) + + def forward(self, X): + """ + Forward pass through the MLP. + + Arguments: + X (torch.Tensor): Output tensor of shape `(batch_size, ...)` + + Returns: + + torch.Tensor: Output of the MLP with shape `(batch_size, linear_sizes[-1])` + """ + return self.net(X) diff --git a/torchmil/nn/variational_autoencoder.py b/torchmil/nn/variational_autoencoder.py new file mode 100644 index 0000000..5b7f32a --- /dev/null +++ b/torchmil/nn/variational_autoencoder.py @@ -0,0 +1,604 @@ +import torch +from .mlp import MLP +import numpy as np + + +class VariationalAutoEncoder(torch.nn.Module): + r""" + Variational Autoencoder (VAE) model for learning latent representations. + + The VAE learns a latent representation $\mathbf{z}$ of input data $\mathbf{x}$ by maximizing the Evidence Lower Bound (ELBO): + + $$ + \mathcal{L}(\theta, \phi; \mathbf{x}) = \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} [\log p_\theta(\mathbf{x}|\mathbf{z})] - \text{KL}(q_\phi(\mathbf{z}|\mathbf{x}) \| p(\mathbf{z})) + $$ + + where $q_\phi(\mathbf{z}|\mathbf{x})$ is the encoder (posterior) and $p_\theta(\mathbf{x}|\mathbf{z})$ is the decoder (likelihood). + Both the encoder and decoder are implemented as MLPs. + """ + + def __init__( + self, + input_shape: tuple[int] = (512,), + layer_sizes: list[int] = [128, 64], + activations: list[str] = ["relu", "None"], + covar_mode: str = "single", + jitter: float = 1e-7, + ) -> None: + """ + Arguments: + input_shape: Shape of input data (excluding batch dimension). + layer_sizes: List of hidden layer sizes for the encoder (decoder mirrors this). + activations: List of activation functions for each layer. Must have same length as layer_sizes. + covar_mode: Covariance mode for the variational distributions. Options: 'single', 'diagonal'. + jitter: Small value added to log_std for numerical stability. + """ + super().__init__() + + self.input_dim = input_shape + self.output_size = input_shape[0] + self.jitter = jitter + self.covar_mode = covar_mode + self.layer_sizes = layer_sizes + + dimensions_enc = [input_shape[0]] + layer_sizes + dimensions_dec = layer_sizes[::-1] + [input_shape[0]] + + # Compute Variance dimensions + if covar_mode == "single": + self.d_var_enc, self.d_var_dec = 1, 1 + elif covar_mode == "diagonal": + self.d_var_enc, self.d_var_dec = layer_sizes[-1], input_shape[0] + else: + raise NotImplementedError( + f"{covar_mode} covar mode not valid. Current implementations: single/diagonal" + ) + + dimensions_enc[-1] += self.d_var_enc + dimensions_dec[-1] += self.d_var_dec + + self.encoder = MLP( + input_size=dimensions_enc[0], + linear_sizes=dimensions_enc[1:], + activations=["relu" for i in range(len(dimensions_enc) - 2)] + ["None"], + ) + self.decoder = MLP( + input_size=dimensions_dec[0], + linear_sizes=dimensions_dec[1:], + activations=["relu" for i in range(len(dimensions_dec) - 2)] + ["None"], + ) + + def get_reparameterized_samples( + self, mean: torch.Tensor, log_std: torch.Tensor, n_samples: int = 1 + ) -> torch.Tensor: + """ + Generate reparameterized samples using the reparameterization trick. + + Arguments: + mean: Mean of the distribution of shape `(batch_size, latent_dim)`. + log_std: Log standard deviation of shape `(batch_size, d_var_enc)`. + n_samples: Number of samples to generate. + + Returns: + Reparameterized samples of shape `(batch_size, n_samples, latent_dim)`. + """ + + # Obtain samples + samples = torch.normal( + 0.0, 1.0, size=(mean.shape[0], mean.shape[-1], n_samples) + ).to( + mean.device + ) # (batch_size, latent_dim, n_samples) + + # Reparameterized samples, add jitter + log_std = log_std + self.jitter + rep_samples = samples * torch.exp(log_std).unsqueeze(-1) + mean.unsqueeze( + -1 + ) # (batch_size, latent_dim, n_samples) + rep_samples = rep_samples.transpose(1, 2) # (batch_size, n_samples, latent_dim) + + return rep_samples + + def get_raw_output_enc(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute the mean and log standard deviation of the posterior distribution $q(\mathbf{z}\mid \mathbf{x})$. + + The posterior distribution is parameterized as: + + $q(\mathbf{z} \mid \mathbf{x}) = \mathcal N(\mathbf{x} \mid \mu(\mathbf{x}), \sigma(\mathbf{x}) * \mathbf{I})$ + + Arguments: + X: Input data of shape `(batch_size, input_dim)`. + + Returns: + mean: Mean vector of shape `(batch_size, latent_dim)`. + log_std: Log standard deviation of shape `(batch_size, d_var_enc)`. + """ + # Reshape in case of images + if len(X.shape) > 3: + X = torch.flatten(X, start_dim=1) # (batch_size, input_dim) + + out = self.encoder(X) # (batch_size, latent_dim + d_var_enc) + + mean, log_std = ( + out[:, : -self.d_var_enc], + out[:, -self.d_var_enc :], + ) # (batch_size, latent_dim), (batch_size, d_var_enc) + return mean, log_std + + def get_raw_output_dec( + self, samples: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute the mean and log standard deviation of the likelihood distribution $p(\mathbf{x}|\mathbf{z})$. + + The likelihood distribution is parameterized as: + $p(\mathbf{x} \mid \mathbf{z}) = \mathcal N (\mathbf{x} \mid \mu(\mathbf{z}), \sigma(\mathbf{z}) \mathbf{I}) + + Arguments: + samples: Samples from the posterior distribution of shape `(batch_size, latent_dim)`. + + Returns: + mean: Mean of the likelihood of shape `(batch_size, input_dim)`. + log_std: Log standard deviation of shape `(batch_size, d_var_dec)`. + """ + + # Decode Samples + out_rec = self.decoder(samples) # (batch_size, input_dim + d_var_dec) + # Obtain mean and var + mean, log_std = ( + out_rec[:, : -self.d_var_dec], + out_rec[:, -self.d_var_dec :], + ) # (batch_size, input_dim), (batch_size, d_var_dec) + + if self.covar_mode == "single": + log_std = torch.ones_like(mean) * log_std + + return mean, log_std + + def forward( + self, X: torch.Tensor, n_samples: int = 1, return_mean_logstd: bool = False + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Forward pass through the VAE encoder. + + Note: This method only implements encoding, since the latent variables are used for downstream tasks. + + Arguments: + X: Input data of shape `(batch_size, ...)`. + n_samples: Number of Monte Carlo samples to generate from the posterior. + return_mean_logstd: If True, also returns the posterior mean and log standard deviation. + + Returns: + posterior_samples: Samples from the posterior $q(\mathbf{z}|\mathbf{x})$ of shape `(batch_size, n_samples, latent_dim)`. + post_mean: Only returned when `return_mean_logstd=True`. Posterior mean of shape `(batch_size, latent_dim)`. + post_log_std: Only returned when `return_mean_logstd=True`. Posterior log std of shape `(batch_size, latent_dim)`. + """ + + if len(X.shape) > 3: + X = torch.flatten(X, start_dim=1) + + posterior_samples, post_mean, post_log_std = self.get_posterior_samples( + X=X, n_samples=n_samples, return_mean_logstd=True + ) + + if return_mean_logstd: + return posterior_samples, post_mean, post_log_std + + return posterior_samples + + def get_posterior_samples( + self, X: torch.Tensor, n_samples: int = 1, return_mean_logstd: bool = False + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Generate samples from the posterior distribution $q(\mathbf{z}|\mathbf{x})$. + + Arguments: + X: Input data of shape `(batch_size, input_dim)`. + n_samples: Number of samples to obtain. + return_mean_logstd: Whether to return the mean and log_std used to obtain the samples. + + Returns: + rep_samples: Samples of $q(\mathbf{z}|\mathbf{x})$ of shape `(batch_size, n_samples, latent_dim)`. + mean: Only returned when `return_mean_logstd=True`. Mean of shape `(batch_size, latent_dim)`. + log_std_v: Only returned when `return_mean_logstd=True`. Log std of shape `(batch_size, latent_dim)`. + """ + # Flatten in case of images + if len(X.shape) > 2: + X = torch.flatten(X, start_dim=1) + + mean, log_std = self.get_raw_output_enc( + X + ) # (batch_size, latent_dim), (batch_size, d_var_enc) + + # Current implementation only stands single variance for all the latent vector + if self.covar_mode == "single": + log_std_v = torch.ones_like(mean) * log_std # (batch_size, latent_dim) + elif self.covar_mode == "diagonal": + log_std_v = log_std # (batch_size, latent_dim) + + rep_samples = self.get_reparameterized_samples( + mean, log_std_v, n_samples=n_samples + ) # (batch_size, n_samples, latent_dim) + + if return_mean_logstd: + return rep_samples, mean, log_std_v + + return rep_samples + + def complete_forward_samples( + self, X: torch.Tensor, n_samples: int = 1 + ) -> torch.Tensor: + """ + Compute samples from the likelihood $p(\mathbf{x}|\mathbf{z})$ using a complete forward pass. + + This method first encodes the input to obtain latent samples, then decodes + these samples to reconstruct the original input. + + Arguments: + X: Input data of shape `(batch_size, input_dim)`. + n_samples: Number of Monte Carlo samples for the forward pass. + + Returns: + reconstructions: Reconstructed data of shape `(batch_size, input_dim)`. + """ + orig_shape = X.shape + + samples, mean, log_std = self.get_posterior_samples( + X, n_samples=n_samples, return_mean_logstd=True + ) # (batch_size, n_samples, latent_dim), (batch_size, latent_dim), (batch_size, latent_dim) + + # Flatten the MC dimension + samples = samples.reshape( + samples.shape[0] * samples.shape[1], samples.shape[2] + ) # (batch_size * n_samples, latent_dim) + # Compute likelihood mean and log_std + lik_mean, lik_log_std = self.get_raw_output_dec( + samples + ) # (batch_size * n_samples, input_dim), (batch_size * n_samples, d_var_dec) + + # Reshape back to (batch_size, n_samples, input_dim) and take mean over samples + lik_mean = lik_mean.view( + orig_shape[0], n_samples, orig_shape[1] + ) # (batch_size, n_samples, input_dim) + lik_mean = lik_mean.mean( + dim=1 + ) # (batch_size, input_dim) - average over samples + + return lik_mean + + def compute_loss( + self, + X: torch.Tensor, + reduction: str = "sum", + n_samples: int = 1, + return_samples: bool = False, + ) -> dict | tuple[dict, torch.Tensor]: + r""" + Compute the ELBO: + + $\mathcal E_q[log p(\mathbf{x}|\mathbf{z})] - KL[q(\mathbf{z})||p(\mathbf{z})]$ + + Arguments: + X: Input data of shape `(batch_size, input_dim)`. + reduction: Way to reduce the loss across instances. Options: 'sum', 'mean', 'none'. + n_samples: Number of Monte Carlo samples for the loss computation. + return_samples: If True, also returns the latent samples used for loss computation. + + Returns: + loss_dict: Dictionary containing 'VaeELL' (negative expected log-likelihood) and 'VaeKL' (KL divergence). + samples: Only returned when `return_samples=True`. Latent samples of shape `(batch_size * n_samples, latent_dim)`. + """ + + if len(X.shape) > 2: + X = torch.flatten(X, start_dim=1) # (batch_size, ...) + + samples, post_mean, post_log_std = self.get_posterior_samples( + X, n_samples=n_samples, return_mean_logstd=True + ) # (batch_size, n_samples, latent_dim), (batch_size, latent_dim), (batch_size, latent_dim) + + # Flatten the MC dimension + samples = samples.reshape( + samples.shape[0] * samples.shape[1], samples.shape[2] + ) # (batch_size * n_samples, latent_dim) + + # Compute likelihood mean and log_std + lik_mean, lik_log_std = self.get_raw_output_dec( + samples + ) # (batch_size * n_samples, input_dim), (batch_size * n_samples, d_var_dec) + + # Replicate for each MC sample + X_replicated = X.repeat_interleave(n_samples, dim=0) + + # Compute for all inputs LL + ell = self._diagonal_log_gaussian_pdf( + X_replicated, lik_mean, lik_log_std + ) # (batch_size * n_samples) + + kl = self._kl_prior(post_mean, post_log_std) # (batch_size) + + # Reshape the ELL and compute mean in MC samples + ell = ell.view(X.shape[0], n_samples).mean( + dim=1 + ) # (batch_size, n_samples) -> (batch_size) + + # Reduce dimensions, change sign to KL + if reduction == "mean": + ell = torch.mean(ell) # () + kl = torch.mean(kl) # () + elif reduction == "sum": + ell = torch.sum(ell) # () + kl = torch.sum(kl) # () + + # Care: Current implementation returns -ELL + if return_samples: + return {"VaeELL": -ell, "VaeKL": kl}, samples + return {"VaeELL": -ell, "VaeKL": kl} + + def _kl_prior(self, mean: torch.Tensor, log_std: torch.Tensor) -> torch.Tensor: + """ + Compute KL divergence between posterior $q(\mathbf{z}|\mathbf{x})$ and standard normal prior. + + Computes $D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}) || \mathcal{N}(0, I))$ for a multivariate Gaussian + posterior with diagonal covariance matrix. + + Arguments: + mean: Posterior mean vectors of shape `(batch_size, latent_dim)`. + log_std: Posterior log standard deviations of shape `(batch_size, latent_dim)`. + + Returns: + kl_div: KL divergence per instance of shape `(batch_size,)`. + """ + + kl_div = -0.5 * torch.sum( + 1 + 2 * (log_std) - mean**2 - torch.exp(2 * log_std), dim=1 + ) + + return kl_div + + def _diagonal_log_gaussian_pdf( + self, inputs: torch.Tensor, mean: torch.Tensor, log_std: torch.Tensor + ) -> torch.Tensor: + r""" + Compute log probability density of a diagonal Gaussian. + + Computes $\log \mathcal{N}(x; \mu, \sigma^2 I)$ for inputs with diagonal covariance. + + Arguments: + inputs: Input data of shape `(batch_size, input_dim)`. + mean: Gaussian mean of shape `(batch_size, input_dim)`. + log_std: Gaussian log standard deviation of shape `(batch_size, input_dim)`. + + Returns: + log_prob: Log probability densities of shape `(batch_size,)`. + """ + + # Const + K = inputs.shape[-1] + log_std = log_std + + # Const term. Det of Sigma is prod of diag. + log_det_cov = 2 * torch.sum(log_std, dim=1) + + # Compute terms + inv_cov = torch.exp(-2 * log_std) + diff = inputs - mean + + quadratic_term = torch.sum(diff * diff * inv_cov, dim=1) + + # log density + log_pdf = -0.5 * ( + K * torch.log(torch.tensor(2 * np.pi)) + log_det_cov + quadratic_term + ) + + return log_pdf + + def log_marginal_likelihood_importance_sampling( + self, X: torch.Tensor, n_samples: int = 1 + ) -> torch.Tensor: + r""" + Compute log marginal likelihood log $p(\mathbf{x})$ via importance sampling. The estimation is computed as + + $\log p(\mathbf{x}) \approx \log \frac{1}{K} \sum_{i=1}^K \frac{p(x|z_i)p(z_i)}{q(z_i|x)}$ + + Arguments: + X: Input data of shape `(batch_size, input_dim)`. + n_samples: Number of importance samples for estimation. + + Returns: + log_marginal: Log marginal likelihood estimates of shape `(batch_size,)`. + """ + # Case of bag of images + if len(X.shape) > 2: + X = torch.flatten(X, start_dim=1) # (batch_size, ...) + + # Encode + samples, mean_post, log_std_post = self.get_posterior_samples( + X, n_samples=n_samples, return_mean_logstd=True + ) # (batch_size, n_samples, latent_dim), (batch_size, latent_dim), (batch_size, latent_dim) + samples = samples.reshape( + samples.shape[0] * samples.shape[1], -1 + ) # (batch_size * n_samples, latent_dim) + + # Decode Samples + mean_lik, log_std_lik = self.get_raw_output_dec( + samples + ) # (batch_size * n_samples, input_dim), (batch_size * n_samples, d_var_dec) + + # Replicate for each MC sample + X_replicated = X.repeat_interleave( + n_samples, dim=0 + ) # (batch_size * n_samples, input_dim) + mean_post = mean_post.repeat_interleave( + n_samples, dim=0 + ) # (batch_size * n_samples, latent_dim) + log_std_post = log_std_post.repeat_interleave( + n_samples, dim=0 + ) # (batch_size * n_samples, latent_dim) + + # Compute marginal LLs + ll = self._diagonal_log_gaussian_pdf(X_replicated, mean_lik, log_std_lik) + prior_log_lik = self._diagonal_log_gaussian_pdf( + samples, torch.zeros_like(samples), torch.ones_like(samples) + ) + post_log_lik = self._diagonal_log_gaussian_pdf(samples, mean_post, log_std_post) + + # Compute via log sum exp as: + # log p(x) = log (1/n_samples) * sum_i exp(log p(x|z_i) + log p(z_i) - log q(z_i|x)) + log_exponent = ll + prior_log_lik - post_log_lik # (batch_size * n_samples) + log_exponent = log_exponent.view( + X.shape[0], n_samples + ) # (batch_size, n_samples) + unnormalized_log_marginal_imp_sampling = torch.logsumexp( + log_exponent, dim=1 + ) # (batch_size) + log_marginal_importance = unnormalized_log_marginal_imp_sampling + torch.log( + torch.tensor(1 / n_samples, dtype=torch.float64, device=X.device) + ) + + return log_marginal_importance + + +class VariationalAutoEncoderMIL(VariationalAutoEncoder): + r""" + Variational Autoencoder for Multiple Instance Learning. + + This class extends the VAE to handle bag-structured data by processing each instance + in a bag independently through the VAE and returning results in bag format. + """ + + def forward( + self, X: torch.Tensor, n_samples: int = 1, return_mean_logstd: bool = False + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Forward pass for bag-structured data. + + This method processes each instance in the bags independently through the VAE encoder. + Used in MIL feature extraction where output must be `(batch_size, bag_size, latent_dim)`. + + Arguments: + X: Bag data of shape `(batch_size, bag_size, input_dim)`. + n_samples: Number of Monte Carlo samples. + return_mean_logstd: Whether to return posterior mean and log standard deviation. + + Returns: + samples: Encoding samples of shape `(batch_size, bag_size, n_samples, latent_dim)`. + mean: Only returned when `return_mean_logstd=True`. Mean of shape `(batch_size, bag_size, latent_dim)`. + log_std: Only returned when `return_mean_logstd=True`. Log std of shape `(batch_size, bag_size, latent_dim)`. + """ + if len(X.shape) == 2: + X = X.unsqueeze(0) + + B, N = X.shape[0], X.shape[1] + + # View as individual elements and forward + samples, mean, log_std = super().forward( + X.view(B * N, *X.shape[2:]), n_samples, return_mean_logstd=True + ) # (batch_size * bag_size, n_samples, latent_dim), # (batch_size * bag_size, latent_dim), (batch_size * bag_size, latent_dim) + + samples = samples.view( + B, N, n_samples, mean.shape[-1] + ) # (batch_size, bag_size, n_samples, latent_dim) + + if return_mean_logstd: + mean = mean.view(B, N, mean.shape[-1]) # (batch_size, bag_size, latent_dim) + log_std = log_std.view( + B, N, log_std.shape[-1] + ) # (batch_size, bag_size, latent_dim) + return samples, mean, log_std + + return samples + + def log_marginal_likelihood_importance_sampling( + self, X: torch.Tensor, mask: torch.Tensor | None = None, n_samples: int = 1 + ) -> torch.Tensor: + """ + Compute log marginal likelihood for bag-structured data via importance sampling. + + This method processes each instance in the bags independently and returns + log marginal estimates for each instance. + + Arguments: + X: Bag data of shape `(batch_size, bag_size, input_dim)`. + mask: Optional binary mask of shape `(batch_size, bag_size)` indicating valid instances. + n_samples: Number of importance samples for estimation. + + Returns: + log_marginal: Log marginal likelihood per instance of shape `(batch_size, bag_size)`. + """ + B, N = X.shape[0], X.shape[1] + mask = mask if mask is not None else torch.ones(B, N).to(X.device) + log_marginal_imp = super().log_marginal_likelihood_importance_sampling( + X.view(B * N, *X.shape[2:]), n_samples + ) # (batch_size * bag_size) + return log_marginal_imp.view(B, N) * mask + + def compute_loss( + self, + X: torch.Tensor, + mask: torch.Tensor | None = None, + reduction: str = "mean", + n_samples: int = 1, + return_samples: bool = False, + ) -> dict | tuple[dict, torch.Tensor]: + """ + Compute VAE loss for bag-structured data. + + The loss is computed for each instance in the bags and then aggregated according + to the reduction strategy and optional mask. + + Arguments: + X: Bag data of shape `(batch_size, bag_size, input_dim)`. + mask: Optional binary mask of shape `(batch_size, bag_size)` for valid instances. + reduction: Reduction method ('sum', 'mean', or 'none'). + n_samples: Number of Monte Carlo samples for loss computation. + return_samples: Whether to return latent samples used in loss computation. + + Returns: + loss_dict: Dictionary with 'VaeELL' and 'VaeKL' losses. + samples: Only returned when `return_samples=True`. Latent samples used for loss computation. + """ + + B, N, D = X.shape[0], X.shape[1], X.shape[2:] + mask = mask if mask is not None else torch.ones(B, N) + + X = X.view(B * N, *X.shape[2:]) + + # Recall that Super returns {-ell, KL} + losses, samples = super().compute_loss( + X, reduction="none", n_samples=n_samples, return_samples=True + ) + + # Assumes that the MC samples are reduced in the VAE + losses["VaeELL"] = losses["VaeELL"].view(B, N) + losses["VaeKL"] = losses["VaeKL"].view(B, N) + + if reduction == "sum": + losses["VaeELL"] = torch.sum(losses["VaeELL"].sum(dim=-1)) + losses["VaeKL"] = torch.sum(losses["VaeKL"].sum(dim=-1)) + elif reduction == "mean": + losses["VaeELL"] = torch.mean(losses["VaeELL"].mean(dim=-1)) + losses["VaeKL"] = torch.mean(losses["VaeKL"].mean(dim=-1)) + + if return_samples: + return losses, samples.view(B, n_samples, N, -1) + return losses + + def complete_forward_samples(self, X: torch.Tensor) -> torch.Tensor: + """ + Compute reconstructions for bag-structured data via complete forward pass. + + This method processes each instance in the bags independently through the + VAE encoder-decoder pipeline and returns reconstructions in bag format. + + Arguments: + X: Bag data of shape `(batch_size, bag_size, input_dim)`. + + Returns: + reconstructions: Reconstructed bag data of shape `(batch_size, bag_size, input_dim)`. + """ + B, N, D = X.shape[0], X.shape[1], *X.shape[2:] + recs = super().complete_forward_samples( + X.view(B * N, *X.shape[2:]) + ) # (batch_size * bag_size * n_samples) + return recs.view(B, N, D) From 36c5f505df7d12dbdc719c8b617ce0e740e178be Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 5 Dec 2025 23:29:21 +0100 Subject: [PATCH 02/10] Removed unused variable --- torchmil/nn/variational_autoencoder.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/torchmil/nn/variational_autoencoder.py b/torchmil/nn/variational_autoencoder.py index 5b7f32a..84c9a9c 100644 --- a/torchmil/nn/variational_autoencoder.py +++ b/torchmil/nn/variational_autoencoder.py @@ -100,7 +100,7 @@ def get_reparameterized_samples( return rep_samples def get_raw_output_enc(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """ + r""" Compute the mean and log standard deviation of the posterior distribution $q(\mathbf{z}\mid \mathbf{x})$. The posterior distribution is parameterized as: @@ -129,7 +129,7 @@ def get_raw_output_enc(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso def get_raw_output_dec( self, samples: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """ + r""" Compute the mean and log standard deviation of the likelihood distribution $p(\mathbf{x}|\mathbf{z})$. The likelihood distribution is parameterized as: @@ -159,7 +159,7 @@ def get_raw_output_dec( def forward( self, X: torch.Tensor, n_samples: int = 1, return_mean_logstd: bool = False ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ + r""" Forward pass through the VAE encoder. Note: This method only implements encoding, since the latent variables are used for downstream tasks. @@ -190,7 +190,7 @@ def forward( def get_posterior_samples( self, X: torch.Tensor, n_samples: int = 1, return_mean_logstd: bool = False ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ + r""" Generate samples from the posterior distribution $q(\mathbf{z}|\mathbf{x})$. Arguments: @@ -229,7 +229,7 @@ def get_posterior_samples( def complete_forward_samples( self, X: torch.Tensor, n_samples: int = 1 ) -> torch.Tensor: - """ + r""" Compute samples from the likelihood $p(\mathbf{x}|\mathbf{z})$ using a complete forward pass. This method first encodes the input to obtain latent samples, then decodes @@ -336,7 +336,7 @@ def compute_loss( return {"VaeELL": -ell, "VaeKL": kl} def _kl_prior(self, mean: torch.Tensor, log_std: torch.Tensor) -> torch.Tensor: - """ + r""" Compute KL divergence between posterior $q(\mathbf{z}|\mathbf{x})$ and standard normal prior. Computes $D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}) || \mathcal{N}(0, I))$ for a multivariate Gaussian @@ -559,7 +559,7 @@ def compute_loss( samples: Only returned when `return_samples=True`. Latent samples used for loss computation. """ - B, N, D = X.shape[0], X.shape[1], X.shape[2:] + B, N, _ = X.shape[0], X.shape[1], X.shape[2:] mask = mask if mask is not None else torch.ones(B, N) X = X.view(B * N, *X.shape[2:]) From 78d5217c08cae2545250c9332a7257d790751011 Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 5 Dec 2025 23:33:53 +0100 Subject: [PATCH 03/10] Reformatting ruff --- tests/models/test_vaeabmil.py | 25 ++++++--------- tests/nn/test_mlp.py | 16 +++++----- tests/nn/test_variational_autoencoder.py | 40 +++++++++--------------- torchmil/models/__init__.py | 2 +- torchmil/nn/__init__.py | 6 ++-- torchmil/nn/variational_autoencoder.py | 4 +-- 6 files changed, 40 insertions(+), 53 deletions(-) diff --git a/tests/models/test_vaeabmil.py b/tests/models/test_vaeabmil.py index cde7b63..3f89efc 100644 --- a/tests/models/test_vaeabmil.py +++ b/tests/models/test_vaeabmil.py @@ -20,19 +20,14 @@ def sample_data(): def vae_feat_ext(): # Returns a VariationalAutoEncoderMIL instance for feature extraction return VariationalAutoEncoderMIL( - input_shape=(10,), - layer_sizes=[8, 5], - activations=["relu", "None"] + input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] ) @pytest.fixture def vaeabmil_model(vae_feat_ext): # Returns an instance of the VAEABMIL model with default parameters - return VAEABMIL( - feat_ext=vae_feat_ext, - in_shape=(3, 10) - ) + return VAEABMIL(feat_ext=vae_feat_ext, in_shape=(3, 10)) # Basic tests for VAEABMIL class @@ -45,7 +40,7 @@ def test_vaeabmil_initialization(vae_feat_ext): def test_vaeabmil_forward_pass(sample_data, vaeabmil_model): # Test basic forward pass X, _, mask = sample_data - + Y_pred = vaeabmil_model(X, mask) assert Y_pred.shape == (2,) @@ -53,7 +48,7 @@ def test_vaeabmil_forward_pass(sample_data, vaeabmil_model): def test_vaeabmil_forward_with_attention(sample_data, vaeabmil_model): # Test forward pass with attention return X, _, mask = sample_data - + Y_pred, att = vaeabmil_model(X, mask, return_att=True) assert Y_pred.shape == (2,) assert att.shape == (2, 3) @@ -62,13 +57,13 @@ def test_vaeabmil_forward_with_attention(sample_data, vaeabmil_model): def test_vaeabmil_compute_loss(sample_data, vaeabmil_model): # Test loss computation X, Y, mask = sample_data - + # Ensure the model is in a good state for loss computation with torch.no_grad(): _ = vaeabmil_model(X, mask) - + Y_pred, loss_dict = vaeabmil_model.compute_loss(Y, X, mask) - + assert Y_pred.shape == (2,) assert "BCEWithLogitsLoss" in loss_dict assert "VaeELL" in loss_dict @@ -78,10 +73,10 @@ def test_vaeabmil_compute_loss(sample_data, vaeabmil_model): def test_vaeabmil_predict(sample_data, vaeabmil_model): # Test predict method X, _, mask = sample_data - + Y_pred = vaeabmil_model.predict(X, mask, return_inst_pred=False) assert Y_pred.shape == (2,) - + Y_pred, y_inst_pred = vaeabmil_model.predict(X, mask, return_inst_pred=True) assert Y_pred.shape == (2,) - assert y_inst_pred.shape == (2, 3) \ No newline at end of file + assert y_inst_pred.shape == (2, 3) diff --git a/tests/nn/test_mlp.py b/tests/nn/test_mlp.py index a8d46cc..d34e7dc 100644 --- a/tests/nn/test_mlp.py +++ b/tests/nn/test_mlp.py @@ -9,13 +9,13 @@ def test_get_activation(): # Test different activation functions relu = get_activation("relu") assert isinstance(relu, torch.nn.ReLU) - + sigmoid = get_activation("sigmoid") assert isinstance(sigmoid, torch.nn.Sigmoid) - + tanh = get_activation("tanh") assert isinstance(tanh, torch.nn.Tanh) - + # Test None case none_activation = get_activation("none") assert none_activation is None @@ -66,12 +66,12 @@ def test_mlp_no_activation(): def test_mlp_different_activations(): # Test different activation functions work input_data = torch.randn(3, 10) - + # ReLU activation mlp_relu = MLP(input_size=10, linear_sizes=[5], activations=["relu"]) output_relu = mlp_relu(input_data) assert output_relu.shape == (3, 5) - + # Sigmoid activation mlp_sigmoid = MLP(input_size=10, linear_sizes=[5], activations=["sigmoid"]) output_sigmoid = mlp_sigmoid(input_data) @@ -80,7 +80,9 @@ def test_mlp_different_activations(): def test_mlp_multiple_layers(): # Test MLP with multiple layers - mlp = MLP(input_size=20, linear_sizes=[16, 8, 4], activations=["relu", "relu", "None"]) + mlp = MLP( + input_size=20, linear_sizes=[16, 8, 4], activations=["relu", "relu", "None"] + ) input_data = torch.randn(2, 20) output = mlp(input_data) - assert output.shape == (2, 4) \ No newline at end of file + assert output.shape == (2, 4) diff --git a/tests/nn/test_variational_autoencoder.py b/tests/nn/test_variational_autoencoder.py index 32e6fb4..a141174 100644 --- a/tests/nn/test_variational_autoencoder.py +++ b/tests/nn/test_variational_autoencoder.py @@ -1,7 +1,10 @@ import torch import pytest -from torchmil.nn.variational_autoencoder import VariationalAutoEncoder, VariationalAutoEncoderMIL +from torchmil.nn.variational_autoencoder import ( + VariationalAutoEncoder, + VariationalAutoEncoderMIL, +) # Fixtures for common setup @@ -18,18 +21,14 @@ def sample_bag_data(): @pytest.fixture def vae_basic(): return VariationalAutoEncoder( - input_shape=(10,), - layer_sizes=[8, 5], - activations=["relu", "None"] + input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] ) @pytest.fixture def vae_mil(): return VariationalAutoEncoderMIL( - input_shape=(10,), - layer_sizes=[8, 5], - activations=["relu", "None"] + input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] ) @@ -37,9 +36,7 @@ def vae_mil(): def test_vae_initialization(): # Test basic initialization vae = VariationalAutoEncoder( - input_shape=(10,), - layer_sizes=[8, 5], - activations=["relu", "None"] + input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] ) assert vae.input_dim == (10,) assert vae.output_size == 10 @@ -49,9 +46,7 @@ def test_vae_initialization(): def test_vae_initialization_diagonal_covar(): # Test initialization with diagonal covariance vae = VariationalAutoEncoder( - input_shape=(10,), - layer_sizes=[8, 5], - covar_mode="diagonal" + input_shape=(10,), layer_sizes=[8, 5], covar_mode="diagonal" ) assert vae.covar_mode == "diagonal" @@ -60,9 +55,7 @@ def test_vae_initialization_invalid_covar(): # Test that invalid covariance mode raises error with pytest.raises(NotImplementedError): VariationalAutoEncoder( - input_shape=(10,), - layer_sizes=[8, 5], - covar_mode="invalid" + input_shape=(10,), layer_sizes=[8, 5], covar_mode="invalid" ) @@ -87,7 +80,7 @@ def test_vae_complete_forward_samples(sample_data, vae_basic): def test_vae_compute_loss(sample_data, vae_basic): # Test loss computation loss_dict = vae_basic.compute_loss(sample_data, reduction="sum", n_samples=2) - + assert "VaeELL" in loss_dict assert "VaeKL" in loss_dict assert loss_dict["VaeELL"].shape == () # scalar @@ -97,7 +90,7 @@ def test_vae_compute_loss(sample_data, vae_basic): def test_vae_get_raw_output_enc(sample_data, vae_basic): # Test encoder raw output mean, log_std = vae_basic.get_raw_output_enc(sample_data) - + assert mean.shape == (4, 5) # batch_size, latent_dim assert log_std.shape == (4, 1) # batch_size, d_var_enc (single mode) @@ -106,7 +99,7 @@ def test_vae_get_raw_output_dec(vae_basic): # Test decoder raw output latent_samples = torch.randn(4, 5) # batch_size, latent_dim mean, log_std = vae_basic.get_raw_output_dec(latent_samples) - + assert mean.shape == (4, 10) # batch_size, input_dim # In single covar mode, log_std is expanded to match input dim assert log_std.shape == (4, 10) # batch_size, input_dim (expanded from d_var_dec) @@ -115,10 +108,7 @@ def test_vae_get_raw_output_dec(vae_basic): # Basic tests for VariationalAutoEncoderMIL class def test_vae_mil_initialization(): # Test MIL VAE initialization - vae_mil = VariationalAutoEncoderMIL( - input_shape=(10,), - layer_sizes=[8, 5] - ) + vae_mil = VariationalAutoEncoderMIL(input_shape=(10,), layer_sizes=[8, 5]) assert isinstance(vae_mil, VariationalAutoEncoder) @@ -131,7 +121,7 @@ def test_vae_mil_forward(sample_bag_data, vae_mil): def test_vae_mil_compute_loss(sample_bag_data, vae_mil): # Test MIL VAE loss computation loss_dict = vae_mil.compute_loss(sample_bag_data, reduction="mean") - + assert "VaeELL" in loss_dict and "VaeKL" in loss_dict assert loss_dict["VaeELL"].shape == () assert loss_dict["VaeKL"].shape == () @@ -140,4 +130,4 @@ def test_vae_mil_compute_loss(sample_bag_data, vae_mil): def test_vae_mil_complete_forward_samples(sample_bag_data, vae_mil): # Test complete forward pass for MIL VAE reconstructions = vae_mil.complete_forward_samples(sample_bag_data) - assert reconstructions.shape == sample_bag_data.shape \ No newline at end of file + assert reconstructions.shape == sample_bag_data.shape diff --git a/torchmil/models/__init__.py b/torchmil/models/__init__.py index 84a0e9f..7e3cea7 100644 --- a/torchmil/models/__init__.py +++ b/torchmil/models/__init__.py @@ -26,4 +26,4 @@ from .iibmil import IIBMIL as IIBMIL from .setmil import SETMIL as SETMIL from .gtp import GTP as GTP -from .vaeabmil import VAEABMIL as VAEABMIL \ No newline at end of file +from .vaeabmil import VAEABMIL as VAEABMIL diff --git a/torchmil/nn/__init__.py b/torchmil/nn/__init__.py index 1a377df..b90ae2a 100644 --- a/torchmil/nn/__init__.py +++ b/torchmil/nn/__init__.py @@ -41,5 +41,7 @@ ) from .mlp import MLP as MLP -from .variational_autoencoder import (VariationalAutoEncoder as VariationalAutoEncoder, - VariationalAutoEncoderMIL as VariationalAutoEncoderMIL) \ No newline at end of file +from .variational_autoencoder import ( + VariationalAutoEncoder as VariationalAutoEncoder, + VariationalAutoEncoderMIL as VariationalAutoEncoderMIL, +) diff --git a/torchmil/nn/variational_autoencoder.py b/torchmil/nn/variational_autoencoder.py index 84c9a9c..b404ce9 100644 --- a/torchmil/nn/variational_autoencoder.py +++ b/torchmil/nn/variational_autoencoder.py @@ -86,9 +86,7 @@ def get_reparameterized_samples( # Obtain samples samples = torch.normal( 0.0, 1.0, size=(mean.shape[0], mean.shape[-1], n_samples) - ).to( - mean.device - ) # (batch_size, latent_dim, n_samples) + ).to(mean.device) # (batch_size, latent_dim, n_samples) # Reparameterized samples, add jitter log_std = log_std + self.jitter From 13202e951408781246d1dabd9b63f95de6858378 Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 6 Feb 2026 19:00:03 +0100 Subject: [PATCH 04/10] Adding Video Classification Dataset and TAD Dataset. --- tests/datasets/test_tad_dataset.py | 93 +++++++++++++ .../test_video_classification_dataset.py | 69 ++++++++++ torchmil/datasets/__init__.py | 2 + torchmil/datasets/tadmil_dataset.py | 92 +++++++++++++ .../datasets/video_classification_dataset.py | 127 ++++++++++++++++++ torchmil/nn/variational_autoencoder.py | 6 +- 6 files changed, 387 insertions(+), 2 deletions(-) create mode 100644 tests/datasets/test_tad_dataset.py create mode 100644 tests/datasets/test_video_classification_dataset.py create mode 100644 torchmil/datasets/tadmil_dataset.py create mode 100644 torchmil/datasets/video_classification_dataset.py diff --git a/tests/datasets/test_tad_dataset.py b/tests/datasets/test_tad_dataset.py new file mode 100644 index 0000000..87915e6 --- /dev/null +++ b/tests/datasets/test_tad_dataset.py @@ -0,0 +1,93 @@ +import pytest +import numpy as np +import pandas as pd +from torchmil.datasets import TADMILDataset + +@pytest.fixture +def mock_tad_dataset(tmp_path): + root = tmp_path + features = "resnet50" + features_path = root / f"features/features_{features}" + labels_path = root / "labels" + frame_labels_path = root / "frame_labels" + + # Create directories + for path in [features_path, labels_path, frame_labels_path]: + path.mkdir(parents=True, exist_ok=True) + + # Create a dummy video + video_name = "video1" + num_frames = 5 + feature_dim = 128 + + # Save dummy .npy files + # Features: (5 frames, 128 dimensions) + np.save(features_path / f"{video_name}.npy", np.random.rand(num_frames, feature_dim)) + # Label: Scalar (0 or 1) + np.save(labels_path / f"{video_name}.npy", np.array(1)) + # Frame Labels: (5 frames,) + np.save(frame_labels_path / f"{video_name}.npy", np.random.randint(0, 2, size=(num_frames,))) + + # Create a splits.csv + splits = pd.DataFrame({"bag_name": [video_name], "split": ["train"]}) + splits.to_csv(root / "splits.csv", index=False) + + return { + "root": str(root), + "features": features, + "partition": "train", + "bag_keys": ["X", "Y", "y_inst", "coords"], # Explicitly requesting y_inst for test + "adj_with_dist": False, + "norm_adj": True, + "load_at_init": False, + } + +def test_tad_init(mock_tad_dataset): + dataset = TADMILDataset(**mock_tad_dataset) + assert hasattr(dataset, "_load_bag") + # Verify the path construction logic inside init + assert dataset.features_path.endswith("features/features_resnet50/") + assert dataset.labels_path.endswith("labels/") + assert dataset.inst_labels_path.endswith("frame_labels/") + +def test_tad_load_bag(mock_tad_dataset): + dataset = TADMILDataset(**mock_tad_dataset) + bag = dataset._load_bag("video1") + + assert isinstance(bag, dict) + assert "X" in bag + assert "Y" in bag + assert "y_inst" in bag # Should be present because we added it to bag_keys in fixture + assert "coords" in bag # Added automatically by _add_coords + + # Verify shapes match the created data + assert bag["X"].shape[0] == 5 # 5 frames + assert bag["coords"].shape[0] == 5 + assert bag["y_inst"].shape[0] == 5 + +def test_tad_split_filtering(mock_tad_dataset, tmp_path): + """Test that it respects the partition (train vs test) defined in splits.csv""" + + # Add a 'test' video to the existing structure + root = tmp_path + features_path = root / "features/features_resnet50" + + video_test = "video_test" + np.save(features_path / f"{video_test}.npy", np.random.rand(5, 128)) + + # Update splits.csv to include a test video + splits = pd.DataFrame({ + "bag_name": ["video1", video_test], + "split": ["train", "test"] + }) + splits.to_csv(root / "splits.csv", index=False) + + # Initialize with partition='test' + params = mock_tad_dataset.copy() + params["partition"] = "test" + + dataset = TADMILDataset(**params) + + # Should only contain 'video_test' + assert len(dataset.bag_names) == 1 + assert dataset.bag_names[0] == "video_test" \ No newline at end of file diff --git a/tests/datasets/test_video_classification_dataset.py b/tests/datasets/test_video_classification_dataset.py new file mode 100644 index 0000000..6dd7fab --- /dev/null +++ b/tests/datasets/test_video_classification_dataset.py @@ -0,0 +1,69 @@ +import pytest +import numpy as np +from torchmil.datasets import VideoClassificationDataset + +@pytest.fixture +def mock_video_data(tmp_path): + features_path = tmp_path / "features" + labels_path = tmp_path / "labels" + frame_labels_path = tmp_path / "frame_labels" + + for p in [features_path, labels_path, frame_labels_path]: + p.mkdir(parents=True, exist_ok=True) + + video_name = "vid1" + num_frames = 5 + feature_dim = 64 + + # Save dummy .npy files + # Features: (5 frames, 64 features) + np.save(features_path / f"{video_name}.npy", np.random.rand(num_frames, feature_dim)) + # Label: Scalar + np.save(labels_path / f"{video_name}.npy", np.array(0)) + # Frame Labels: (5 frames,) + np.save(frame_labels_path / f"{video_name}.npy", np.random.randint(0, 2, size=(num_frames,))) + + return { + "features_path": str(features_path), + "labels_path": str(labels_path), + "frame_labels_path": str(frame_labels_path), + "video_names": [video_name], + "bag_keys": ["X", "Y", "coords", "y_inst"], + "adj_with_dist": False, + "norm_adj": True, + "load_at_init": False, + } + +def test_video_dataset_init(mock_video_data): + dataset = VideoClassificationDataset(**mock_video_data) + # Check that hardcoded defaults in __init__ are set correctly + assert dataset.dist_thr == 1.10 + # Check that arguments were mapped correctly (e.g. video_names -> bag_names) + assert dataset.bag_names == ["vid1"] + assert hasattr(dataset, "_load_bag") + +def test_load_bag_coords(mock_video_data): + dataset = VideoClassificationDataset(**mock_video_data) + bag = dataset._load_bag(mock_video_data["video_names"][0]) + + assert "coords" in bag + assert isinstance(bag["coords"], np.ndarray) + + # Coords should be shape (N, 1) + assert bag["coords"].shape == (bag["X"].shape[0], 1) + + # Coords should be sequential integers [0, 1, 2, 3, 4] + assert np.array_equal(bag["coords"].flatten(), np.arange(bag["X"].shape[0])) + +def test_load_bag_contents(mock_video_data): + """Test that all requested keys (X, Y, y_inst) are loaded correctly.""" + dataset = VideoClassificationDataset(**mock_video_data) + bag = dataset._load_bag(mock_video_data["video_names"][0]) + + assert "X" in bag + assert bag["X"].shape == (5, 64) + + assert "Y" in bag + + assert "y_inst" in bag + assert bag["y_inst"].shape[0] == 5 \ No newline at end of file diff --git a/torchmil/datasets/__init__.py b/torchmil/datasets/__init__.py index c3a3f0d..dadecfc 100644 --- a/torchmil/datasets/__init__.py +++ b/torchmil/datasets/__init__.py @@ -6,10 +6,12 @@ from .camelyon16mil_dataset import CAMELYON16MILDataset as CAMELYON16MILDataset from .pandamil_dataset import PANDAMILDataset as PANDAMILDataset from .rsnamil_dataset import RSNAMILDataset as RSNAMILDataset +from .tadmil_dataset import TADMILDataset as TADMILDataset from .toy_dataset import ToyDataset as ToyDataset from .wsi_dataset import WSIDataset as WSIDataset from .ctscan_dataset import CTScanDataset as CTScanDataset +from .video_classification_dataset import VideoClassificationDataset from .mc_standard_dataset import MCStandardMILDataset as MCStandardMILDataset from .sc_standard_dataset import SCStandardMILDataset as SCStandardMILDataset diff --git a/torchmil/datasets/tadmil_dataset.py b/torchmil/datasets/tadmil_dataset.py new file mode 100644 index 0000000..3b570db --- /dev/null +++ b/torchmil/datasets/tadmil_dataset.py @@ -0,0 +1,92 @@ +import numpy as np + +from .binary_classification_dataset import BinaryClassificationDataset +from .video_classification_dataset import VideoClassificationDataset + +from ..utils.common import read_csv, keep_only_existing_files + + +class TADMILDataset(BinaryClassificationDataset, VideoClassificationDataset): + r""" + Traffic Anomaly Detection for Multiple Instance Learning (MIL). + Download it from [Kaggle Datasets](https://www.kaggle.com/datasets/nikanvasei/traffic-anomaly-dataset-tad). + + + **Dataset description.** + We have preprocessed the Video by computing features for each frame using various feature extractors. + + - A **video** is labeled as positive (`frame_label=1`) if it contains evidence of traffic anomaly. + - A **video** is labeled as positive (`label=1`) if it contains at least one positive frame. + + This means a video is considered positive if there is any evidence of traffic anomaly. + + **Directory structure.** + + The following directory structure is expected: + + ``` + root + ├── features + │ ├── features_{features} + │ │ ├── video1.npy + │ │ ├── video2.npy + │ │ └── ... + ├── labels + │ ├── video1.npy + │ ├── video2.npy + │ └── ... + └── splits.csv + ``` + + Each `.npy` file corresponds to a video. The `splits.csv` file defines train/test splits for standardized experimentation. + """ + + def __init__( + self, + root: str, + features: str = "resnet50", + partition: str = "train", + bag_keys: list = ["X", "Y", "adj", "coords"], + adj_with_dist: bool = False, + norm_adj: bool = True, + load_at_init: bool = True, + ) -> None: + """ + Arguments: + root: Path to the root directory of the dataset. + features: Type of features to use. Must be one of ['resnet18', 'resnet50', 'vit_b_32'] + partition: Partition of the dataset. Must be one of ['train', 'test']. + bag_keys: List of keys to use for the bags. Must be in ['X', 'Y', 'y_inst', 'coords']. + adj_with_dist: If True, the adjacency matrix is built using the Euclidean distance between the patches features. If False, the adjacency matrix is binary. + norm_adj: If True, normalize the adjacency matrix. + load_at_init: If True, load the bags at initialization. If False, load the bags on demand. + """ + features_path = f"{root}/features/features_{features}/" + labels_path = f"{root}/labels/" + frame_labels_path = f"{root}/frame_labels/" + + splits_file = f"{root}/splits.csv" + dict_list = read_csv(splits_file) + video_names = [ + row["bag_name"] for row in dict_list if row["split"] == partition + ] + + video_names = list(set(video_names)) + video_names = keep_only_existing_files(features_path, video_names) + + VideoClassificationDataset.__init__( + self, + features_path=features_path, + labels_path=labels_path, + frame_labels_path=frame_labels_path, + bag_keys=bag_keys, + video_names=video_names, + adj_with_dist=adj_with_dist, + norm_adj=norm_adj, + load_at_init=load_at_init, + ) + + def _load_bag(self, name: str) -> dict[str, np.ndarray]: + bag_dict = BinaryClassificationDataset._load_bag(self, name) + bag_dict = VideoClassificationDataset._add_coords(self, bag_dict) + return bag_dict diff --git a/torchmil/datasets/video_classification_dataset.py b/torchmil/datasets/video_classification_dataset.py new file mode 100644 index 0000000..a9967a1 --- /dev/null +++ b/torchmil/datasets/video_classification_dataset.py @@ -0,0 +1,127 @@ +import numpy as np + +from .processed_mil_dataset import ProcessedMILDataset + + +class VideoClassificationDataset(ProcessedMILDataset): + r""" + This class represents a dataset of videos for Multiple Instance Learning (MIL). + + **MIL and Video Classification.** + Videos are sequences of frames that capture motion and temporal information. + In the context of MIL, a video is considered a bag, and the frames are considered instances. + + **Directory structure.** + It is assumed that the bags have been processed and saved as numpy files. + For more information on the processing of the bags, refer to the [`ProcessedMILDataset` class](processed_mil_dataset.md). + This dataset expects the following directory structure: + + ``` + features_path + ├── video1.npy + ├── video2.npy + └── ... + labels_path + ├── video1.npy + ├── video2.npy + └── ... + inst_labels_path + ├── video1.npy + ├── video2.npy + └── ... + ``` + + **Order of the frames and the adjacency matrix.** + This dataset assumes that the frames of the video frames are ordered. + An adjacency matrix $\mathbf{A} = \left[ A_{ij} \right]$ is built using this information: + + \begin{equation} + A_{ij} = \begin{cases} + d_{ij}, & \text{if } \lvert i - j \rvert = 1, \\ + 0, & \text{otherwise}, + \end{cases} \quad d_{ij} = \begin{cases} + 1, & \text{if } \text{adj_with_dist=False}, \\ + \exp\left( -\frac{\left\| \mathbf{x}_i - \mathbf{x}_j \right\|}{d} \right), & \text{if } \text{adj_with_dist=True}. + \end{cases} + \end{equation} + + where $\mathbf{x}_i \in \mathbb{R}^d$ and $\mathbf{x}_j \in \mathbb{R}^d$ are the features of instances $i$ and $j$, respectively. + """ + + def __init__( + self, + features_path: str, + labels_path: str, + frame_labels_path: str = None, + video_names: list = None, + bag_keys: list = ["X", "Y", "y_inst", "adj", "coords"], + adj_with_dist: bool = False, + norm_adj: bool = True, + load_at_init: bool = True, + ) -> None: + """ + Class constructor. + + Arguments: + features_path: Path to the directory containing the matrices of the videos + labels_path: Path to the directory containing the labels of the videos. + frame_labels_path: Path to the directory containing the labels of the frames. + video_names: List of the names of the videos to load. If None, all videos in the `features_path` directory are loaded. + bag_keys: List of keys to use for the bags. Must be in ['X', 'Y', 'y_inst', 'coords']. + adj_with_dist: If True, the adjacency matrix is built using the Euclidean distance between the frames features. If False, the adjacency matrix is binary. + norm_adj: If True, normalize the adjacency matrix. + load_at_init: If True, load the bags at initialization. If False, load the bags on demand. + """ + + dist_thr = 1.10 + super().__init__( + features_path=features_path, + labels_path=labels_path, + inst_labels_path=frame_labels_path, + coords_path="", + bag_names=video_names, + bag_keys=bag_keys, + adj_with_dist=adj_with_dist, + dist_thr=dist_thr, + norm_adj=norm_adj, + load_at_init=load_at_init, + ) + + def _add_coords(self, bag_dict: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + """ + Add coordinates to the bag dictionary. + + Arguments: + bag_dict: Dictionary containing the features, label and instance labels of the bag. + + Returns: + bag_dict: Dictionary containing the features, label, instance labels and coordinates of the bag. + """ + bag_size = bag_dict["X"].shape[0] + bag_dict["coords"] = np.arange(0, bag_size).reshape(-1, 1) + + return bag_dict + + def _load_bag(self, name: str) -> dict[str, np.ndarray]: + """ + Load a bag from disk. + + Arguments: + name: Name of the bag to load. + + Returns: + bag_dict: Dictionary containing the features, label, instance labels and coordinates of the bag. + """ + bag_dict = {} + if "X" in self.bag_keys: + bag_dict["X"] = self._load_features(name) + + if "Y" in self.bag_keys: + bag_dict["Y"] = self._load_labels(name) + + if "y_inst" in self.bag_keys: + bag_dict["y_inst"] = self._load_inst_labels(name) + + if "coords" in self.bag_keys or "adj" in self.bag_keys: + bag_dict = self._add_coords(bag_dict) + return bag_dict diff --git a/torchmil/nn/variational_autoencoder.py b/torchmil/nn/variational_autoencoder.py index b404ce9..0d37f4c 100644 --- a/torchmil/nn/variational_autoencoder.py +++ b/torchmil/nn/variational_autoencoder.py @@ -86,7 +86,9 @@ def get_reparameterized_samples( # Obtain samples samples = torch.normal( 0.0, 1.0, size=(mean.shape[0], mean.shape[-1], n_samples) - ).to(mean.device) # (batch_size, latent_dim, n_samples) + ).to( + mean.device + ) # (batch_size, latent_dim, n_samples) # Reparameterized samples, add jitter log_std = log_std + self.jitter @@ -268,7 +270,7 @@ def complete_forward_samples( def compute_loss( self, X: torch.Tensor, - reduction: str = "sum", + reduction: str = "mean", n_samples: int = 1, return_samples: bool = False, ) -> dict | tuple[dict, torch.Tensor]: From 9327c81d40943ad59d275b71e6650abcf90f4a99 Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 6 Feb 2026 19:07:07 +0100 Subject: [PATCH 05/10] Ruff format --- tests/datasets/test_tad_dataset.py | 51 ++++++++++++------- .../test_video_classification_dataset.py | 25 ++++++--- torchmil/nn/variational_autoencoder.py | 4 +- 3 files changed, 51 insertions(+), 29 deletions(-) diff --git a/tests/datasets/test_tad_dataset.py b/tests/datasets/test_tad_dataset.py index 87915e6..1b0c866 100644 --- a/tests/datasets/test_tad_dataset.py +++ b/tests/datasets/test_tad_dataset.py @@ -3,6 +3,7 @@ import pandas as pd from torchmil.datasets import TADMILDataset + @pytest.fixture def mock_tad_dataset(tmp_path): root = tmp_path @@ -19,14 +20,19 @@ def mock_tad_dataset(tmp_path): video_name = "video1" num_frames = 5 feature_dim = 128 - + # Save dummy .npy files # Features: (5 frames, 128 dimensions) - np.save(features_path / f"{video_name}.npy", np.random.rand(num_frames, feature_dim)) + np.save( + features_path / f"{video_name}.npy", np.random.rand(num_frames, feature_dim) + ) # Label: Scalar (0 or 1) np.save(labels_path / f"{video_name}.npy", np.array(1)) # Frame Labels: (5 frames,) - np.save(frame_labels_path / f"{video_name}.npy", np.random.randint(0, 2, size=(num_frames,))) + np.save( + frame_labels_path / f"{video_name}.npy", + np.random.randint(0, 2, size=(num_frames,)), + ) # Create a splits.csv splits = pd.DataFrame({"bag_name": [video_name], "split": ["train"]}) @@ -36,12 +42,18 @@ def mock_tad_dataset(tmp_path): "root": str(root), "features": features, "partition": "train", - "bag_keys": ["X", "Y", "y_inst", "coords"], # Explicitly requesting y_inst for test + "bag_keys": [ + "X", + "Y", + "y_inst", + "coords", + ], # Explicitly requesting y_inst for test "adj_with_dist": False, "norm_adj": True, "load_at_init": False, } + def test_tad_init(mock_tad_dataset): dataset = TADMILDataset(**mock_tad_dataset) assert hasattr(dataset, "_load_bag") @@ -50,6 +62,7 @@ def test_tad_init(mock_tad_dataset): assert dataset.labels_path.endswith("labels/") assert dataset.inst_labels_path.endswith("frame_labels/") + def test_tad_load_bag(mock_tad_dataset): dataset = TADMILDataset(**mock_tad_dataset) bag = dataset._load_bag("video1") @@ -57,37 +70,39 @@ def test_tad_load_bag(mock_tad_dataset): assert isinstance(bag, dict) assert "X" in bag assert "Y" in bag - assert "y_inst" in bag # Should be present because we added it to bag_keys in fixture - assert "coords" in bag # Added automatically by _add_coords - + assert ( + "y_inst" in bag + ) # Should be present because we added it to bag_keys in fixture + assert "coords" in bag # Added automatically by _add_coords + # Verify shapes match the created data - assert bag["X"].shape[0] == 5 # 5 frames + assert bag["X"].shape[0] == 5 # 5 frames assert bag["coords"].shape[0] == 5 assert bag["y_inst"].shape[0] == 5 + def test_tad_split_filtering(mock_tad_dataset, tmp_path): """Test that it respects the partition (train vs test) defined in splits.csv""" - + # Add a 'test' video to the existing structure root = tmp_path features_path = root / "features/features_resnet50" - + video_test = "video_test" np.save(features_path / f"{video_test}.npy", np.random.rand(5, 128)) - + # Update splits.csv to include a test video - splits = pd.DataFrame({ - "bag_name": ["video1", video_test], - "split": ["train", "test"] - }) + splits = pd.DataFrame( + {"bag_name": ["video1", video_test], "split": ["train", "test"]} + ) splits.to_csv(root / "splits.csv", index=False) # Initialize with partition='test' params = mock_tad_dataset.copy() params["partition"] = "test" - + dataset = TADMILDataset(**params) - + # Should only contain 'video_test' assert len(dataset.bag_names) == 1 - assert dataset.bag_names[0] == "video_test" \ No newline at end of file + assert dataset.bag_names[0] == "video_test" diff --git a/tests/datasets/test_video_classification_dataset.py b/tests/datasets/test_video_classification_dataset.py index 6dd7fab..b1a8ea5 100644 --- a/tests/datasets/test_video_classification_dataset.py +++ b/tests/datasets/test_video_classification_dataset.py @@ -2,6 +2,7 @@ import numpy as np from torchmil.datasets import VideoClassificationDataset + @pytest.fixture def mock_video_data(tmp_path): features_path = tmp_path / "features" @@ -14,14 +15,19 @@ def mock_video_data(tmp_path): video_name = "vid1" num_frames = 5 feature_dim = 64 - + # Save dummy .npy files # Features: (5 frames, 64 features) - np.save(features_path / f"{video_name}.npy", np.random.rand(num_frames, feature_dim)) + np.save( + features_path / f"{video_name}.npy", np.random.rand(num_frames, feature_dim) + ) # Label: Scalar np.save(labels_path / f"{video_name}.npy", np.array(0)) # Frame Labels: (5 frames,) - np.save(frame_labels_path / f"{video_name}.npy", np.random.randint(0, 2, size=(num_frames,))) + np.save( + frame_labels_path / f"{video_name}.npy", + np.random.randint(0, 2, size=(num_frames,)), + ) return { "features_path": str(features_path), @@ -34,6 +40,7 @@ def mock_video_data(tmp_path): "load_at_init": False, } + def test_video_dataset_init(mock_video_data): dataset = VideoClassificationDataset(**mock_video_data) # Check that hardcoded defaults in __init__ are set correctly @@ -42,19 +49,21 @@ def test_video_dataset_init(mock_video_data): assert dataset.bag_names == ["vid1"] assert hasattr(dataset, "_load_bag") + def test_load_bag_coords(mock_video_data): dataset = VideoClassificationDataset(**mock_video_data) bag = dataset._load_bag(mock_video_data["video_names"][0]) assert "coords" in bag assert isinstance(bag["coords"], np.ndarray) - + # Coords should be shape (N, 1) assert bag["coords"].shape == (bag["X"].shape[0], 1) - + # Coords should be sequential integers [0, 1, 2, 3, 4] assert np.array_equal(bag["coords"].flatten(), np.arange(bag["X"].shape[0])) + def test_load_bag_contents(mock_video_data): """Test that all requested keys (X, Y, y_inst) are loaded correctly.""" dataset = VideoClassificationDataset(**mock_video_data) @@ -62,8 +71,8 @@ def test_load_bag_contents(mock_video_data): assert "X" in bag assert bag["X"].shape == (5, 64) - + assert "Y" in bag - + assert "y_inst" in bag - assert bag["y_inst"].shape[0] == 5 \ No newline at end of file + assert bag["y_inst"].shape[0] == 5 diff --git a/torchmil/nn/variational_autoencoder.py b/torchmil/nn/variational_autoencoder.py index 0d37f4c..8816fde 100644 --- a/torchmil/nn/variational_autoencoder.py +++ b/torchmil/nn/variational_autoencoder.py @@ -86,9 +86,7 @@ def get_reparameterized_samples( # Obtain samples samples = torch.normal( 0.0, 1.0, size=(mean.shape[0], mean.shape[-1], n_samples) - ).to( - mean.device - ) # (batch_size, latent_dim, n_samples) + ).to(mean.device) # (batch_size, latent_dim, n_samples) # Reparameterized samples, add jitter log_std = log_std + self.jitter From 642a35bbbbfce75775af64b13129f24b324b8eea Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 6 Feb 2026 19:12:45 +0100 Subject: [PATCH 06/10] Removing unused statement --- torchmil/datasets/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/torchmil/datasets/__init__.py b/torchmil/datasets/__init__.py index dadecfc..35d05e2 100644 --- a/torchmil/datasets/__init__.py +++ b/torchmil/datasets/__init__.py @@ -11,7 +11,6 @@ from .toy_dataset import ToyDataset as ToyDataset from .wsi_dataset import WSIDataset as WSIDataset from .ctscan_dataset import CTScanDataset as CTScanDataset -from .video_classification_dataset import VideoClassificationDataset from .mc_standard_dataset import MCStandardMILDataset as MCStandardMILDataset from .sc_standard_dataset import SCStandardMILDataset as SCStandardMILDataset From 05e11733d9f5ad49a04303ce3a5ddeb07096d69e Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 6 Feb 2026 19:17:56 +0100 Subject: [PATCH 07/10] Fixing test --- tests/datasets/test_video_classification_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/datasets/test_video_classification_dataset.py b/tests/datasets/test_video_classification_dataset.py index b1a8ea5..5a73812 100644 --- a/tests/datasets/test_video_classification_dataset.py +++ b/tests/datasets/test_video_classification_dataset.py @@ -1,6 +1,6 @@ import pytest import numpy as np -from torchmil.datasets import VideoClassificationDataset +from torchmil.datasets.video_classification_dataset import VideoClassificationDataset @pytest.fixture From 79d15fa7e0269f983f08049cd068c0c4f03376cf Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 6 Feb 2026 20:11:00 +0100 Subject: [PATCH 08/10] Improving tests --- .../datasets/test_mc_standard_mil_dataset.py | 311 +++++++++++------- tests/datasets/test_trident_wsi_dataset.py | 290 +++++++++------- tests/nn/test_variational_autoencoder.py | 294 +++++++++++------ torchmil/datasets/mc_standard_dataset.py | 32 -- 4 files changed, 549 insertions(+), 378 deletions(-) diff --git a/tests/datasets/test_mc_standard_mil_dataset.py b/tests/datasets/test_mc_standard_mil_dataset.py index 04fba6d..362ba17 100644 --- a/tests/datasets/test_mc_standard_mil_dataset.py +++ b/tests/datasets/test_mc_standard_mil_dataset.py @@ -2,140 +2,201 @@ import torch from torchmil.datasets import MCStandardMILDataset - -def test_mcstandardmilda_init(): - """ - Tests the initialization of MCStandardMILDataset. - """ - dataset = MCStandardMILDataset(D=5, num_bags=10, pos_class_prob=0.7, seed=42) - assert dataset.num_bags == 10 - assert dataset.pos_class_prob == 0.7 - assert dataset.train is True - assert len(dataset) == 10 # Check __len__ - - # Verify distributions are initialized - assert len(dataset.pos_distr) == 2 - assert dataset.neg_distr is not None - assert dataset.poisoning is not None - - -def test_mcstandardmilda_len(): - """ - Tests the __len__ method of MCStandardMILDataset. - """ - dataset = MCStandardMILDataset(D=2, num_bags=50) - assert len(dataset) == 50 - +# --- Fixtures --- + +@pytest.fixture +def default_dataset(): + """Returns a standard dataset for general testing.""" + # We seed torch here to ensure the fixture itself is reproducible across test runs + torch.manual_seed(42) + return MCStandardMILDataset(D=5, num_bags=10, pos_class_prob=0.5, seed=42) + +# --- Tests --- + +def test_init_properties(default_dataset): + """Test initialization of attributes and distributions.""" + ds = default_dataset + assert ds.num_bags == 10 + assert ds.pos_class_prob == 0.5 + assert ds.train is True + assert len(ds) == 10 + + # Verify distributions exist + assert isinstance(ds.pos_distr, list) + assert len(ds.pos_distr) == 2 + assert isinstance(ds.neg_distr, torch.distributions.Normal) + assert isinstance(ds.poisoning, torch.distributions.Normal) + +def test_len_empty_edge_case(): + """Test edge case for dataset length.""" dataset_empty = MCStandardMILDataset(D=2, num_bags=0) assert len(dataset_empty) == 0 +def test_getitem_bounds(default_dataset): + """Test __getitem__ boundary conditions.""" + ds = default_dataset + num_bags = len(ds) + + # Valid positive index + _ = ds[0] + _ = ds[num_bags - 1] + + # Valid negative index + _ = ds[-1] + + # Invalid positive index (Explicit raise in code) + with pytest.raises(IndexError, match="out of range"): + _ = ds[num_bags] + + # Invalid negative index (Implicit raise by list) + with pytest.raises(IndexError): + _ = ds[-(num_bags + 1)] -def test_mcstandardmilda_getitem(): +def test_getitem_structure_and_shapes(): """ - Tests the __getitem__ method and bag structure for MCStandardMILDataset. + Test that X (features) and y_inst (instance labels) have matching dimensions. + This covers the stack/view/cat logic in the sample methods. """ - D_val = 3 - num_bags_val = 20 - dataset = MCStandardMILDataset(D=D_val, num_bags=num_bags_val, seed=1) - - # Test valid index - bag = dataset[0] - assert isinstance(bag, dict) - assert "X" in bag - assert "Y" in bag - assert "y_inst" in bag - - assert bag["X"].shape[1] == D_val # Feature dimensionality - assert bag["Y"].ndim == 0 # Scalar label - assert bag["y_inst"].ndim == 1 # 1D instance labels - - # Test out-of-bounds index - with pytest.raises(IndexError): - dataset[num_bags_val] - with pytest.raises(IndexError): - dataset[-num_bags_val - 1] # Test negative index out of bounds - - -def test_mcstandardmilda_positive_bag_content(): + D = 4 + # Seeding torch to ensure sampling doesn't hit edge cases (though unlikely here) + torch.manual_seed(123) + ds = MCStandardMILDataset(D=D, num_bags=5, seed=123) + + for i in range(len(ds)): + bag = ds[i] + X = bag["X"] + y_inst = bag["y_inst"] + Y = bag["Y"] + + # Check types + assert isinstance(X, torch.Tensor) + assert isinstance(y_inst, torch.Tensor) + + # Check dimensions + assert X.ndim == 2 + assert X.shape[1] == D + assert y_inst.ndim == 1 + + # CRITICAL: Number of instances must match + assert X.shape[0] == y_inst.shape[0] + + # Check label consistency + if Y.item() == 1: + # Positive bags must have positive instances (label 1) + assert (y_inst == 1).sum() > 0 + +def test_determinism_and_seeding(): """ - Tests the content of a positive bag in training and test mode. + Test that the dataset generation is reproducible. + + NOTE: The dataset implementation only seeds NumPy internally. + It uses PyTorch for sampling, so we must manually seed PyTorch + in the test to guarantee identical bags. """ - D_val = 2 - # Train mode: no poisoning instance - dataset_train = MCStandardMILDataset( - D=D_val, num_bags=1, pos_class_prob=1.0, train=True, seed=10 - ) - pos_bag_train = dataset_train[0] - assert pos_bag_train["Y"].item() == 1 # Bag label is positive - # Check if poisoning instance (-1) is NOT present in train mode for positive bags - assert -1 not in pos_bag_train["y_inst"] - assert torch.any(pos_bag_train["y_inst"] == 1) # Must have positive instances - assert torch.any(pos_bag_train["y_inst"] == 0) # Must have negative instances - - # Test mode: poisoning instance present - dataset_test = MCStandardMILDataset( - D=D_val, num_bags=1, pos_class_prob=1.0, train=False, seed=10 - ) - pos_bag_test = dataset_test[0] - assert pos_bag_test["Y"].item() == 1 # Bag label is positive - # Check if poisoning instance (-1) IS present in test mode for positive bags - assert -1 in pos_bag_test["y_inst"] - assert torch.any(pos_bag_test["y_inst"] == 1) # Must have positive instances - assert torch.any(pos_bag_test["y_inst"] == 0) # Must have negative instances - - -def test_mcstandardmilda_negative_bag_content(): + seed = 999 + + # Run 1 + torch.manual_seed(seed) + ds1 = MCStandardMILDataset(D=3, num_bags=10, seed=seed) + + # Run 2 + torch.manual_seed(seed) + ds2 = MCStandardMILDataset(D=3, num_bags=10, seed=seed) + + # Run 3 (Control: different seed) + torch.manual_seed(123) + ds3 = MCStandardMILDataset(D=3, num_bags=10, seed=123) + + # Check Exact Match + for i in range(10): + # We check both the data content and the instance labels + assert torch.equal(ds1[i]["X"], ds2[i]["X"]), f"Bag {i} data mismatch" + assert torch.equal(ds1[i]["y_inst"], ds2[i]["y_inst"]), f"Bag {i} labels mismatch" + assert ds1[i]["Y"] == ds2[i]["Y"] + + # Check Mismatch (Sanity check that seeding actually works) + # The first bag is highly likely to differ + assert not torch.equal(ds1[0]["X"], ds3[0]["X"]) + +def test_train_vs_test_poisoning_logic(): """ - Tests the content of a negative bag in training and test mode. + Strictly verify the poisoning logic: + - Train: Negative bags have poison (label -1). + - Test: Positive bags have poison (label -1). """ - D_val = 2 - # Train mode: poisoning instance present - dataset_train = MCStandardMILDataset( - D=D_val, num_bags=1, pos_class_prob=0.0, train=True, seed=11 - ) - neg_bag_train = dataset_train[0] - assert neg_bag_train["Y"].item() == 0 # Bag label is negative - # Check if poisoning instance (-1) IS present in train mode for negative bags - assert -1 in neg_bag_train["y_inst"] - assert torch.any( - neg_bag_train["y_inst"] == 1 - ) # Must have positive instances (single) - assert torch.any(neg_bag_train["y_inst"] == 0) # Must have negative instances - - # Test mode: no poisoning instance - dataset_test = MCStandardMILDataset( - D=D_val, num_bags=1, pos_class_prob=0.0, train=False, seed=11 - ) - neg_bag_test = dataset_test[0] - assert neg_bag_test["Y"].item() == 0 # Bag label is negative - # Check if poisoning instance (-1) is NOT present in test mode for negative bags - assert -1 not in neg_bag_test["y_inst"] - assert torch.any( - neg_bag_test["y_inst"] == 1 - ) # Must have positive instances (single) - assert torch.any(neg_bag_test["y_inst"] == 0) # Must have negative instances - - -def test_mcstandardmilda_bag_counts(): + D = 2 + torch.manual_seed(1) + + # --- Train Mode --- + ds_train = MCStandardMILDataset(D=D, num_bags=20, train=True, seed=1) + + for i in range(len(ds_train)): + bag = ds_train[i] + labels = bag["y_inst"] + is_positive_bag = bag["Y"].item() == 1 + + if is_positive_bag: + # Train Positive: No poison + assert -1 not in labels + else: + # Train Negative: Has poison + assert -1 in labels + # Verify poison values: Mean -10.0 + poison_indices = (labels == -1).nonzero(as_tuple=True)[0] + poison_data = bag["X"][poison_indices] + # Check values are roughly around -10 (far from 0 or 2) + assert torch.all(poison_data < -5.0) + + # --- Test Mode --- + # Re-seed to ensure consistent generation behavior + torch.manual_seed(1) + ds_test = MCStandardMILDataset(D=D, num_bags=20, train=False, seed=1) + + for i in range(len(ds_test)): + bag = ds_test[i] + labels = bag["y_inst"] + is_positive_bag = bag["Y"].item() == 1 + + if is_positive_bag: + # Test Positive: Has poison + assert -1 in labels + poison_indices = (labels == -1).nonzero(as_tuple=True)[0] + poison_data = bag["X"][poison_indices] + assert torch.all(poison_data < -5.0) + else: + # Test Negative: No poison + assert -1 not in labels + +def test_positive_concept_distribution_logic(): """ - Tests that the correct number of positive and negative bags are created. + Verify that positive bags contain data from the positive distributions. + Positive means are 2.0 and 3.0. """ - D_val = 2 - num_bags_total = 100 - pos_prob = 0.6 - dataset = MCStandardMILDataset( - D=D_val, num_bags=num_bags_total, pos_class_prob=pos_prob, seed=12 - ) - - expected_pos_bags = int(num_bags_total * pos_prob) - expected_neg_bags = num_bags_total - expected_pos_bags - - actual_pos_bags = sum( - 1 for i in range(num_bags_total) if dataset[i]["Y"].item() == 1 - ) - actual_neg_bags = sum( - 1 for i in range(num_bags_total) if dataset[i]["Y"].item() == 0 - ) - - assert actual_pos_bags == expected_pos_bags - assert actual_neg_bags == expected_neg_bags + D = 1 + torch.manual_seed(55) + # Create a positive bag in Train mode (to avoid poison noise) + ds = MCStandardMILDataset(D=D, num_bags=1, pos_class_prob=1.0, train=True, seed=55) + bag = ds[0] + + X = bag["X"] + y_inst = bag["y_inst"] + + # Filter for positive instances (label 1) + pos_instances = X[y_inst == 1] + + # Ensure values are strictly positive and reasonably close to means 2.0/3.0 + # (Checking > 1.0 safely excludes the 0.0 negatives and -10.0 poisons) + assert torch.all(pos_instances > 1.0) + assert torch.all(pos_instances < 5.0) + +def test_bag_class_probability(): + """Verify that pos_class_prob controls the class balance.""" + torch.manual_seed(10) + + ds_all_pos = MCStandardMILDataset(D=2, num_bags=10, pos_class_prob=1.0) + labels_all_pos = [ds_all_pos[i]["Y"].item() for i in range(10)] + assert all(label == 1 for label in labels_all_pos) + + ds_all_neg = MCStandardMILDataset(D=2, num_bags=10, pos_class_prob=0.0) + labels_all_neg = [ds_all_neg[i]["Y"].item() for i in range(10)] + assert all(label == 0 for label in labels_all_neg) \ No newline at end of file diff --git a/tests/datasets/test_trident_wsi_dataset.py b/tests/datasets/test_trident_wsi_dataset.py index 6c6a8a0..ee0fb7b 100644 --- a/tests/datasets/test_trident_wsi_dataset.py +++ b/tests/datasets/test_trident_wsi_dataset.py @@ -1,129 +1,183 @@ import pytest import numpy as np -import os -import h5py -from pathlib import Path -from torchmil.datasets import TridentWSIDataset # Update to actual import path +import pandas as pd +from unittest.mock import MagicMock, patch +from torchmil.datasets import TridentWSIDataset -# --- Helper Functions for Fixtures (Kept for creating H5 files) --- - - -def create_h5_file(filepath: Path, dataset_name: str, data: np.ndarray): - """Helper to create a simple HDF5 file with one dataset.""" - filepath.parent.mkdir(parents=True, exist_ok=True) - with h5py.File(filepath, "w") as f: - f.create_dataset(dataset_name, data=data) - - -# --- Pytest Fixtures --- - +# --- Fixtures --- @pytest.fixture -def mock_trident_data(tmp_path): - """ - Sets up a minimal TRIDENT directory structure using H5 files. - """ - WSI_NAME = "sample" - MAG, PS, OPX = 20, 512, 0 - FEAT_EXT = "conch_v15" - TRIDENT_FOLDER = f"{MAG}x_{PS}px_{OPX}px_overlap" - - base_path = tmp_path / "trident_base" - full_trident_path = base_path / TRIDENT_FOLDER - - features_dir = full_trident_path / f"features_{FEAT_EXT}" - labels_dir = full_trident_path / "labels" - coords_dir = full_trident_path / "patches" - inst_labels_dir = full_trident_path / "patch_labels" - - # Create directories - for d in [features_dir, labels_dir, coords_dir, inst_labels_dir]: - d.mkdir(parents=True, exist_ok=True) - - # Raw coordinates (must be divisible by PS=512 for clean test) - # Scaled and normalized result: [[2, 4], [3, 5], [1, 2]] - [1, 2] = [[1, 2], [2, 3], [0, 0]] - raw_coords = np.array([[1024, 2048], [1536, 2560], [512, 1024]]).astype(np.int32) - - # 1. Features file - create_h5_file( - features_dir / f"{WSI_NAME}.h5", - "features", - np.random.rand(raw_coords.shape[0], 128), - ) - - # 2. Label file (WSI-level) - create_h5_file( - labels_dir / f"{WSI_NAME}.h5", "label", np.array([1], dtype=np.float32) - ) - - # 3. Coords file - create_h5_file(coords_dir / f"{WSI_NAME}_patches.h5", "coords", raw_coords) - - # 4. Patch Label file (Instance-level) - create_h5_file( - inst_labels_dir / f"{WSI_NAME}.h5", - "patch_label", - np.random.randint(0, 2, size=(raw_coords.shape[0],)), - ) - +def base_kwargs(): + """Standard arguments for initializing the dataset.""" return { - "base_path": str(base_path) + os.sep, - "labels_path": str(labels_dir) + os.sep, - "patch_labels_path": str(inst_labels_dir) + os.sep, - "wsi_names": [WSI_NAME], - "patch_size": PS, - # The other params are defaults, but included for clarity in TridentWSIDataset - "feature_extractor": FEAT_EXT, - "magnification": MAG, - "overlap_pixels": OPX, - "adj_with_dist": False, - "norm_adj": True, - "load_at_init": False, - "expected_coords": np.array([[1, 2], [2, 3], [0, 0]]).astype(np.int32), + "base_path": "/mock/base", + "labels_path": "/mock/labels", + "feature_extractor": "conch_v15", + "magnification": 20, + "patch_size": 512, + "overlap_pixels": 0, + "wsi_names": ["slide_1", "slide_2"] } - -def test_trident_dataset_init(mock_trident_data): +# --- Tests --- + +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_init_defaults(mock_super_init, base_kwargs): + ds = TridentWSIDataset(**base_kwargs) + + assert ds.patch_size == 512 + assert ds.magnification == 20 + assert ds.overlap_pixels == 0 + assert ds.trident_folder == "20x_512px_0px_overlap/" + + expected_features_path = "/mock/base20x_512px_0px_overlap/features_conch_v15/" + expected_coords_path = "/mock/base20x_512px_0px_overlap/patches/" + + mock_super_init.assert_called_once() + call_kwargs = mock_super_init.call_args[1] + + assert call_kwargs["features_path"] == expected_features_path + assert call_kwargs["coords_path"] == expected_coords_path + assert call_kwargs["dist_thr"] == pytest.approx(np.sqrt(2.0)) + +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_init_custom_threshold(mock_super_init, base_kwargs): + kwargs = base_kwargs.copy() + kwargs["dist_thr"] = 5.5 + + # Fix: Assign to '_' to silence the "unused variable" error + # while still ensuring TridentWSIDataset initializes. + _ = TridentWSIDataset(**kwargs) + + call_kwargs = mock_super_init.call_args[1] + assert call_kwargs["dist_thr"] == 5.5 + +@patch("torchmil.datasets.wsi_dataset.WSIDataset._load_labels") +@patch("os.path.isdir") +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_load_labels_directory_mode(mock_super_init, mock_isdir, mock_super_load, base_kwargs): + """Test directory mode logic. Manually set labels_path since super().__init__ is mocked.""" + mock_isdir.return_value = True + + ds = TridentWSIDataset(**base_kwargs) + ds.labels_path = base_kwargs["labels_path"] # Manually set attribute usually set by super() + + expected_label = np.array([0]) + mock_super_load.return_value = expected_label + + result = ds._load_labels("slide_1") + + assert result == expected_label + mock_isdir.assert_called_with("/mock/labels") + mock_super_load.assert_called_with("slide_1") + +@patch("pandas.read_csv") +@patch("os.path.isdir") +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_load_labels_csv_mode_success(mock_super_init, mock_isdir, mock_read_csv, base_kwargs): + """Test CSV mode logic. Manually set labels_path.""" + mock_isdir.return_value = False + + df = pd.DataFrame({ + "filename": ["slide_1", "slide_2"], + "grade": [0, 1] + }) + mock_read_csv.return_value = df + + kwargs = base_kwargs.copy() + kwargs["wsi_name_col"] = "filename" + kwargs["wsi_label_col"] = "grade" + + ds = TridentWSIDataset(**kwargs) + ds.labels_path = kwargs["labels_path"] # Manually set attribute + + label = ds._load_labels("slide_2") + + assert label[0] == 1 + mock_read_csv.assert_called_once_with("/mock/labels") + +@patch("pandas.read_csv") +@patch("os.path.isdir") +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_load_labels_csv_missing_kwargs(mock_super_init, mock_isdir, mock_read_csv, base_kwargs): + mock_isdir.return_value = False + mock_read_csv.return_value = pd.DataFrame() + + ds = TridentWSIDataset(**base_kwargs) + ds.labels_path = base_kwargs["labels_path"] # Manually set attribute + + with pytest.raises(ValueError, match="must provide 'wsi_name_col' and 'wsi_label_col'"): + ds._load_labels("slide_1") + +@patch("pandas.read_csv") +@patch("os.path.isdir") +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_load_labels_csv_not_found(mock_super_init, mock_isdir, mock_read_csv, base_kwargs): + mock_isdir.return_value = False + + # Mock DF to raise ValueError when accessing specific data + mock_df = MagicMock() + mock_df.loc.__getitem__.side_effect = ValueError("Forced error") + mock_read_csv.return_value = mock_df + + kwargs = base_kwargs.copy() + kwargs["wsi_name_col"] = "name" + kwargs["wsi_label_col"] = "label" + + ds = TridentWSIDataset(**kwargs) + ds.labels_path = kwargs["labels_path"] # Manually set attribute + + with pytest.raises(ValueError, match="Could not read the label"): + ds._load_labels("slide_X") + +@patch("h5py.File") +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_load_coords_calculation(mock_super_init, mock_h5, base_kwargs): + ds = TridentWSIDataset(**base_kwargs) + ds.coords_path = "/mock/patches/" + ds.file_type = ".h5" + ds.patch_size = 512 + + raw_coords = np.array([ + [1024, 2048], + [1536, 2560] + ]) + + mock_file = MagicMock() + mock_file.__getitem__.return_value = raw_coords + mock_h5.return_value = mock_file + + expected_coords = np.array([[0, 0], [1, 1]]) + + result = ds._load_coords("slide_1") + + mock_h5.assert_called_with("/mock/patches/slide_1_patches.h5", "r") + np.testing.assert_array_equal(result, expected_coords) + +@patch("h5py.File") +@patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) +def test_load_coords_none(mock_super_init, mock_h5, base_kwargs): """ - Tests initialization of TridentWSIDataset to check if core attributes are set correctly. + Test _load_coords when the H5 file slice returns None. + Structure: h5py.File(...)['coords'][:] -> None """ - # Assuming TridentWSIDataset is imported or defined above - dataset = TridentWSIDataset(**mock_trident_data) - - # Check attributes set by TridentWSIDataset's __init__ - assert dataset.patch_size == 512 - assert "conch_v15" in dataset.feature_extractor - - # Check attributes passed to super() and derived paths - assert ( - "trident_base/20x_512px_0px_overlap/features_conch_v15/" - in dataset.features_path - ) - assert "trident_base/20x_512px_0px_overlap/patches/" in dataset.coords_path - assert dataset.bag_names == ["sample"] - - -def test_trident_load_coords_adjustment(mock_trident_data): - """ - Tests the overridden _load_coords method for correct scaling and normalization. - """ - # Assuming TridentWSIDataset is imported or defined above - dataset = TridentWSIDataset(**mock_trident_data) - bag_name = mock_trident_data["wsi_names"][0] - - # Check for the presence of the method we expect to be defined - assert hasattr(dataset, "_load_coords") - - loaded_coords = dataset._load_coords(bag_name) - - assert loaded_coords is not None - assert isinstance(loaded_coords, np.ndarray) - print(loaded_coords.dtype) - - # Check normalization and casting as per Trident's logic - assert loaded_coords.min() == 0, "Coordinates were not normalized (min-subtracted)." - assert loaded_coords.dtype == np.int_, "Coordinates should be cast to integer." - assert np.array_equal( - loaded_coords, mock_trident_data["expected_coords"] - ), "Coordinate calculation is incorrect." + ds = TridentWSIDataset(**base_kwargs) + ds.coords_path = "/mock/patches/" + ds.file_type = ".h5" + + # 1. Mock the File object + mock_file_obj = MagicMock() + + # 2. Mock the dataset object returned by file['coords'] + mock_dataset = MagicMock() + + # 3. Mock the slice operator [:] on the dataset to return None + mock_dataset.__getitem__.return_value = None + + # 4. Connect them + mock_file_obj.__getitem__.return_value = mock_dataset + mock_h5.return_value = mock_file_obj + + result = ds._load_coords("slide_empty") + assert result is None \ No newline at end of file diff --git a/tests/nn/test_variational_autoencoder.py b/tests/nn/test_variational_autoencoder.py index a141174..b386fac 100644 --- a/tests/nn/test_variational_autoencoder.py +++ b/tests/nn/test_variational_autoencoder.py @@ -6,128 +6,216 @@ VariationalAutoEncoderMIL, ) +# --- Fixtures --- -# Fixtures for common setup @pytest.fixture -def sample_data(): - return torch.randn(4, 10) # batch_size=4, input_dim=10 +def input_dim(): + return 10 +@pytest.fixture +def latent_dim(): + return 5 @pytest.fixture -def sample_bag_data(): - return torch.randn(2, 3, 10) # batch_size=2, bag_size=3, input_dim=10 +def layer_sizes(latent_dim): + return [8, latent_dim] +@pytest.fixture +def sample_data(input_dim): + return torch.randn(4, input_dim) # batch_size=4, input_dim=10 @pytest.fixture -def vae_basic(): - return VariationalAutoEncoder( - input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] - ) +def sample_image_data(): + # Batch=2, Channels=3, H=4, W=4 -> Flattened size = 48 + return torch.randn(2, 3, 4, 4) +@pytest.fixture +def sample_bag_data(input_dim): + return torch.randn(2, 3, input_dim) # batch_size=2, bag_size=3, input_dim=10 @pytest.fixture -def vae_mil(): - return VariationalAutoEncoderMIL( - input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] +def vae_basic(input_dim, layer_sizes): + return VariationalAutoEncoder( + input_shape=(input_dim,), + layer_sizes=layer_sizes, + activations=["relu", "None"], + covar_mode="single" ) - -# Basic tests for VariationalAutoEncoder class -def test_vae_initialization(): - # Test basic initialization - vae = VariationalAutoEncoder( - input_shape=(10,), layer_sizes=[8, 5], activations=["relu", "None"] +@pytest.fixture +def vae_diagonal(input_dim, layer_sizes): + return VariationalAutoEncoder( + input_shape=(input_dim,), + layer_sizes=layer_sizes, + covar_mode="diagonal" ) - assert vae.input_dim == (10,) - assert vae.output_size == 10 - assert vae.layer_sizes == [8, 5] - -def test_vae_initialization_diagonal_covar(): - # Test initialization with diagonal covariance - vae = VariationalAutoEncoder( - input_shape=(10,), layer_sizes=[8, 5], covar_mode="diagonal" +@pytest.fixture +def vae_mil(input_dim, layer_sizes): + return VariationalAutoEncoderMIL( + input_shape=(input_dim,), + layer_sizes=layer_sizes, + activations=["relu", "None"] ) - assert vae.covar_mode == "diagonal" - - -def test_vae_initialization_invalid_covar(): - # Test that invalid covariance mode raises error - with pytest.raises(NotImplementedError): - VariationalAutoEncoder( - input_shape=(10,), layer_sizes=[8, 5], covar_mode="invalid" - ) - -def test_vae_forward(sample_data, vae_basic): - # Test forward pass (encoding only) +# --- Tests for VariationalAutoEncoder (Standard) --- + +def test_vae_init_validations(): + """Test initialization logic and error handling.""" + # Test valid diagonal init + vae = VariationalAutoEncoder(input_shape=(10,), layer_sizes=[5], covar_mode="diagonal") + assert vae.d_var_enc == 5 # Last layer size + assert vae.d_var_dec == 10 # Input shape + + # Test invalid covar mode + with pytest.raises(NotImplementedError, match="not valid"): + VariationalAutoEncoder(input_shape=(10,), covar_mode="invalid_mode") + +def test_vae_forward_standard(sample_data, vae_basic, latent_dim): + """Test standard forward pass and shape.""" samples = vae_basic(sample_data, n_samples=2) - assert samples.shape == (4, 2, 5) # batch_size, n_samples, latent_dim - - -def test_vae_get_posterior_samples(sample_data, vae_basic): - # Test posterior sampling - samples = vae_basic.get_posterior_samples(sample_data, n_samples=2) - assert samples.shape == (4, 2, 5) # batch_size, n_samples, latent_dim - + assert samples.shape == (4, 2, latent_dim) + +def test_vae_forward_image_input(sample_image_data): + """ + Test the flattening logic in forward/get_posterior_samples. + The input is (B, C, H, W). The VAE must be init with flat dim size. + """ + flat_dim = 3 * 4 * 4 + vae = VariationalAutoEncoder(input_shape=(flat_dim,), layer_sizes=[10]) + + # This hits the `if len(X.shape) > 3` block + samples = vae(sample_image_data, n_samples=1) + assert samples.shape == (2, 1, 10) + +def test_vae_forward_returns_stats(sample_data, vae_basic, latent_dim): + """Test forward with return_mean_logstd=True.""" + samples, mean, log_std = vae_basic(sample_data, n_samples=1, return_mean_logstd=True) + assert samples.shape == (4, 1, latent_dim) + assert mean.shape == (4, latent_dim) + # In 'single' mode, log_std returned by forward is expanded to match mean shape logic? + # Looking at code: `log_std_v = torch.ones_like(mean) * log_std` + assert log_std.shape == (4, latent_dim) + +def test_vae_diagonal_covariance_logic(sample_data, vae_diagonal, input_dim, latent_dim): + """ + Test flow specifically for diagonal covariance mode. + Verifies dimensions of variances in encoder and decoder. + """ + # 1. Encoder Raw Output + mean, log_std = vae_diagonal.get_raw_output_enc(sample_data) + assert mean.shape == (4, latent_dim) + assert log_std.shape == (4, latent_dim) # Diagonal mode: var dim == latent dim + + # 2. Decoder Raw Output + latent_sample = torch.randn(4, latent_dim) + dec_mean, dec_log_std = vae_diagonal.get_raw_output_dec(latent_sample) + assert dec_mean.shape == (4, input_dim) + assert dec_log_std.shape == (4, input_dim) # Diagonal mode: var dim == input dim def test_vae_complete_forward_samples(sample_data, vae_basic): - # Test complete forward pass (encode + decode) - reconstructions = vae_basic.complete_forward_samples(sample_data, n_samples=1) - assert reconstructions.shape == sample_data.shape - - -def test_vae_compute_loss(sample_data, vae_basic): - # Test loss computation - loss_dict = vae_basic.compute_loss(sample_data, reduction="sum", n_samples=2) - + """Test reconstruction path.""" + recs = vae_basic.complete_forward_samples(sample_data, n_samples=5) + # Result is averaged over samples + assert recs.shape == sample_data.shape + +def test_vae_compute_loss_variants(sample_data, vae_basic): + """Test all reduction modes and return flags in compute_loss.""" + # 1. Reduction = Sum + loss_sum = vae_basic.compute_loss(sample_data, reduction="sum") + assert loss_sum["VaeELL"].ndim == 0 + + # 2. Reduction = None (returns per instance) + loss_none = vae_basic.compute_loss(sample_data, reduction="none") + assert loss_none["VaeELL"].shape == (4,) + assert loss_none["VaeKL"].shape == (4,) + + # 3. Return Samples + loss_dict, samples = vae_basic.compute_loss(sample_data, return_samples=True) assert "VaeELL" in loss_dict - assert "VaeKL" in loss_dict - assert loss_dict["VaeELL"].shape == () # scalar - assert loss_dict["VaeKL"].shape == () # scalar - - -def test_vae_get_raw_output_enc(sample_data, vae_basic): - # Test encoder raw output - mean, log_std = vae_basic.get_raw_output_enc(sample_data) - - assert mean.shape == (4, 5) # batch_size, latent_dim - assert log_std.shape == (4, 1) # batch_size, d_var_enc (single mode) - - -def test_vae_get_raw_output_dec(vae_basic): - # Test decoder raw output - latent_samples = torch.randn(4, 5) # batch_size, latent_dim - mean, log_std = vae_basic.get_raw_output_dec(latent_samples) - - assert mean.shape == (4, 10) # batch_size, input_dim - # In single covar mode, log_std is expanded to match input dim - assert log_std.shape == (4, 10) # batch_size, input_dim (expanded from d_var_dec) - - -# Basic tests for VariationalAutoEncoderMIL class -def test_vae_mil_initialization(): - # Test MIL VAE initialization - vae_mil = VariationalAutoEncoderMIL(input_shape=(10,), layer_sizes=[8, 5]) - assert isinstance(vae_mil, VariationalAutoEncoder) - - -def test_vae_mil_forward(sample_bag_data, vae_mil): - # Test MIL VAE forward pass + assert samples.shape[0] == 4 * 1 # Batch * n_samples + +def test_vae_compute_loss_image_flattening(sample_image_data): + """Test that compute_loss handles >2D input (flattening).""" + flat_dim = 3 * 4 * 4 + vae = VariationalAutoEncoder(input_shape=(flat_dim,), layer_sizes=[10]) + # This hits `if len(X.shape) > 2` inside compute_loss + loss = vae.compute_loss(sample_image_data) + assert "VaeELL" in loss + +def test_vae_importance_sampling(sample_data, vae_basic): + """Test log_marginal_likelihood_importance_sampling.""" + log_imp = vae_basic.log_marginal_likelihood_importance_sampling(sample_data, n_samples=10) + assert log_imp.shape == (4,) + + # Test with image data (flattening check) + flat_dim = 3 * 4 * 4 + img_data = torch.randn(2, 3, 4, 4) + vae_img = VariationalAutoEncoder(input_shape=(flat_dim,), layer_sizes=[10]) + log_imp_img = vae_img.log_marginal_likelihood_importance_sampling(img_data, n_samples=2) + assert log_imp_img.shape == (2,) + +# --- Tests for VariationalAutoEncoderMIL (MIL Extension) --- + +def test_mil_forward_structure(sample_bag_data, vae_mil, latent_dim): + """Test basic MIL forward pass dimensions.""" + # Input: (2, 3, 10) -> Output: (2, 3, n_samples, latent) samples = vae_mil(sample_bag_data, n_samples=2) - assert samples.shape == (2, 3, 2, 5) # batch_size, bag_size, n_samples, latent_dim - - -def test_vae_mil_compute_loss(sample_bag_data, vae_mil): - # Test MIL VAE loss computation - loss_dict = vae_mil.compute_loss(sample_bag_data, reduction="mean") - - assert "VaeELL" in loss_dict and "VaeKL" in loss_dict - assert loss_dict["VaeELL"].shape == () - assert loss_dict["VaeKL"].shape == () - - -def test_vae_mil_complete_forward_samples(sample_bag_data, vae_mil): - # Test complete forward pass for MIL VAE - reconstructions = vae_mil.complete_forward_samples(sample_bag_data) - assert reconstructions.shape == sample_bag_data.shape + assert samples.shape == (2, 3, 2, latent_dim) + +def test_mil_forward_single_instance_edge_case(input_dim, vae_mil): + """Test forward pass when input is (BagSize, Dim) instead of (Batch, Bag, Dim).""" + single_bag = torch.randn(5, input_dim) + # The code `if len(X.shape) == 2: X = X.unsqueeze(0)` handles this + samples = vae_mil(single_bag, n_samples=1) + # Output should be (1, 5, 1, latent) + assert samples.shape == (1, 5, 1, vae_mil.layer_sizes[-1]) + +def test_mil_forward_return_stats(sample_bag_data, vae_mil, latent_dim): + """Test return_mean_logstd in MIL context.""" + samples, mean, log_std = vae_mil(sample_bag_data, n_samples=1, return_mean_logstd=True) + assert mean.shape == (2, 3, latent_dim) + assert log_std.shape == (2, 3, latent_dim) + +def test_mil_complete_forward(sample_bag_data, vae_mil): + """Test reconstruction in MIL context.""" + recs = vae_mil.complete_forward_samples(sample_bag_data) + assert recs.shape == sample_bag_data.shape + +def test_mil_compute_loss_masking(sample_bag_data, vae_mil): + """Test loss computation with and without masks, and different reductions.""" + # Mask: 1 for valid, 0 for padding. Let's mask the last instance of bag 0. + mask = torch.ones(2, 3) + mask[0, 2] = 0 + + # 1. Reduction Mean + loss_mean = vae_mil.compute_loss(sample_bag_data, mask=mask, reduction="mean") + assert isinstance(loss_mean["VaeELL"], torch.Tensor) + + # 2. Reduction Sum + loss_sum = vae_mil.compute_loss(sample_bag_data, mask=mask, reduction="sum") + assert isinstance(loss_sum["VaeELL"], torch.Tensor) + + # 3. Reduction None (should return grid) + loss_none = vae_mil.compute_loss(sample_bag_data, mask=mask, reduction="none") + assert loss_none["VaeELL"].shape == (2, 3) + + # 4. Return Samples + loss_dict, samples = vae_mil.compute_loss(sample_bag_data, return_samples=True) + # Expected sample shape: (Batch, n_samples, BagSize, Latent) + # Note: Code returns `samples.view(B, n_samples, N, -1)` + assert samples.shape == (2, 1, 3, vae_mil.layer_sizes[-1]) + +def test_mil_importance_sampling(sample_bag_data, vae_mil): + """Test MIL importance sampling with mask.""" + mask = torch.ones(2, 3) + log_imp = vae_mil.log_marginal_likelihood_importance_sampling( + sample_bag_data, mask=mask, n_samples=5 + ) + assert log_imp.shape == (2, 3) # Returns (Batch, BagSize) + + # Test without mask (defaults to ones) + log_imp_nomask = vae_mil.log_marginal_likelihood_importance_sampling( + sample_bag_data, n_samples=5 + ) + assert log_imp_nomask.shape == (2, 3) \ No newline at end of file diff --git a/torchmil/datasets/mc_standard_dataset.py b/torchmil/datasets/mc_standard_dataset.py index 13cc66a..8dba662 100644 --- a/torchmil/datasets/mc_standard_dataset.py +++ b/torchmil/datasets/mc_standard_dataset.py @@ -163,35 +163,3 @@ def __getitem__(self, index: int) -> TensorDict: f"Index {index} out of range (max: {len(self.bags_list) - 1})" ) return self.bags_list[index] - - -if __name__ == "__main__": - dataset = MCStandardMILDataset(D=2, num_bags=100, pos_class_prob=0.5) - print(f"Number of bags: {len(dataset)}") - for i in range(2): - bag = dataset[i] - print(f"Bag {i}:") - print(f" X: {bag['X']}") - print(f" Y: {bag['Y']}") - print(f" y_inst: {bag['y_inst']}") - bag = dataset[-i] - print(f"Bag {100-i}:") - print(f" X: {bag['X']}") - print(f" Y: {bag['Y']}") - print(f" y_inst: {bag['y_inst']}") - - print("Testing") - dataset_test = MCStandardMILDataset( - D=2, num_bags=100, pos_class_prob=0.5, train=False - ) - for i in range(2): - bag = dataset_test[i] - print(f"Bag {i}:") - print(f" X: {bag['X']}") - print(f" Y: {bag['Y']}") - print(f" y_inst: {bag['y_inst']}") - bag = dataset_test[-i] - print(f"Bag {100-i}:") - print(f" X: {bag['X']}") - print(f" Y: {bag['Y']}") - print(f" y_inst: {bag['y_inst']}") From 8a2776b43123255f22690e01d644dd815bc77d41 Mon Sep 17 00:00:00 2001 From: Francisco Javier Saez Maldonado Date: Fri, 6 Feb 2026 20:18:06 +0100 Subject: [PATCH 09/10] Ruff format --- .../datasets/test_mc_standard_mil_dataset.py | 64 ++++++---- tests/datasets/test_trident_wsi_dataset.py | 112 ++++++++++-------- tests/nn/test_variational_autoencoder.py | 106 +++++++++++------ 3 files changed, 172 insertions(+), 110 deletions(-) diff --git a/tests/datasets/test_mc_standard_mil_dataset.py b/tests/datasets/test_mc_standard_mil_dataset.py index 362ba17..14972fe 100644 --- a/tests/datasets/test_mc_standard_mil_dataset.py +++ b/tests/datasets/test_mc_standard_mil_dataset.py @@ -4,6 +4,7 @@ # --- Fixtures --- + @pytest.fixture def default_dataset(): """Returns a standard dataset for general testing.""" @@ -11,8 +12,10 @@ def default_dataset(): torch.manual_seed(42) return MCStandardMILDataset(D=5, num_bags=10, pos_class_prob=0.5, seed=42) + # --- Tests --- + def test_init_properties(default_dataset): """Test initialization of attributes and distributions.""" ds = default_dataset @@ -20,38 +23,41 @@ def test_init_properties(default_dataset): assert ds.pos_class_prob == 0.5 assert ds.train is True assert len(ds) == 10 - + # Verify distributions exist assert isinstance(ds.pos_distr, list) assert len(ds.pos_distr) == 2 assert isinstance(ds.neg_distr, torch.distributions.Normal) assert isinstance(ds.poisoning, torch.distributions.Normal) + def test_len_empty_edge_case(): """Test edge case for dataset length.""" dataset_empty = MCStandardMILDataset(D=2, num_bags=0) assert len(dataset_empty) == 0 + def test_getitem_bounds(default_dataset): """Test __getitem__ boundary conditions.""" ds = default_dataset num_bags = len(ds) - + # Valid positive index _ = ds[0] _ = ds[num_bags - 1] - + # Valid negative index _ = ds[-1] - + # Invalid positive index (Explicit raise in code) with pytest.raises(IndexError, match="out of range"): _ = ds[num_bags] - + # Invalid negative index (Implicit raise by list) with pytest.raises(IndexError): _ = ds[-(num_bags + 1)] + def test_getitem_structure_and_shapes(): """ Test that X (features) and y_inst (instance labels) have matching dimensions. @@ -61,40 +67,41 @@ def test_getitem_structure_and_shapes(): # Seeding torch to ensure sampling doesn't hit edge cases (though unlikely here) torch.manual_seed(123) ds = MCStandardMILDataset(D=D, num_bags=5, seed=123) - + for i in range(len(ds)): bag = ds[i] X = bag["X"] y_inst = bag["y_inst"] Y = bag["Y"] - + # Check types assert isinstance(X, torch.Tensor) assert isinstance(y_inst, torch.Tensor) - + # Check dimensions assert X.ndim == 2 assert X.shape[1] == D assert y_inst.ndim == 1 - + # CRITICAL: Number of instances must match assert X.shape[0] == y_inst.shape[0] - + # Check label consistency if Y.item() == 1: # Positive bags must have positive instances (label 1) assert (y_inst == 1).sum() > 0 + def test_determinism_and_seeding(): """ Test that the dataset generation is reproducible. - - NOTE: The dataset implementation only seeds NumPy internally. - It uses PyTorch for sampling, so we must manually seed PyTorch + + NOTE: The dataset implementation only seeds NumPy internally. + It uses PyTorch for sampling, so we must manually seed PyTorch in the test to guarantee identical bags. """ seed = 999 - + # Run 1 torch.manual_seed(seed) ds1 = MCStandardMILDataset(D=3, num_bags=10, seed=seed) @@ -111,13 +118,16 @@ def test_determinism_and_seeding(): for i in range(10): # We check both the data content and the instance labels assert torch.equal(ds1[i]["X"], ds2[i]["X"]), f"Bag {i} data mismatch" - assert torch.equal(ds1[i]["y_inst"], ds2[i]["y_inst"]), f"Bag {i} labels mismatch" + assert torch.equal( + ds1[i]["y_inst"], ds2[i]["y_inst"] + ), f"Bag {i} labels mismatch" assert ds1[i]["Y"] == ds2[i]["Y"] # Check Mismatch (Sanity check that seeding actually works) # The first bag is highly likely to differ assert not torch.equal(ds1[0]["X"], ds3[0]["X"]) + def test_train_vs_test_poisoning_logic(): """ Strictly verify the poisoning logic: @@ -126,15 +136,15 @@ def test_train_vs_test_poisoning_logic(): """ D = 2 torch.manual_seed(1) - + # --- Train Mode --- ds_train = MCStandardMILDataset(D=D, num_bags=20, train=True, seed=1) - + for i in range(len(ds_train)): bag = ds_train[i] labels = bag["y_inst"] is_positive_bag = bag["Y"].item() == 1 - + if is_positive_bag: # Train Positive: No poison assert -1 not in labels @@ -145,18 +155,18 @@ def test_train_vs_test_poisoning_logic(): poison_indices = (labels == -1).nonzero(as_tuple=True)[0] poison_data = bag["X"][poison_indices] # Check values are roughly around -10 (far from 0 or 2) - assert torch.all(poison_data < -5.0) + assert torch.all(poison_data < -5.0) # --- Test Mode --- # Re-seed to ensure consistent generation behavior torch.manual_seed(1) ds_test = MCStandardMILDataset(D=D, num_bags=20, train=False, seed=1) - + for i in range(len(ds_test)): bag = ds_test[i] labels = bag["y_inst"] is_positive_bag = bag["Y"].item() == 1 - + if is_positive_bag: # Test Positive: Has poison assert -1 in labels @@ -167,6 +177,7 @@ def test_train_vs_test_poisoning_logic(): # Test Negative: No poison assert -1 not in labels + def test_positive_concept_distribution_logic(): """ Verify that positive bags contain data from the positive distributions. @@ -177,26 +188,27 @@ def test_positive_concept_distribution_logic(): # Create a positive bag in Train mode (to avoid poison noise) ds = MCStandardMILDataset(D=D, num_bags=1, pos_class_prob=1.0, train=True, seed=55) bag = ds[0] - + X = bag["X"] y_inst = bag["y_inst"] - + # Filter for positive instances (label 1) pos_instances = X[y_inst == 1] - + # Ensure values are strictly positive and reasonably close to means 2.0/3.0 # (Checking > 1.0 safely excludes the 0.0 negatives and -10.0 poisons) assert torch.all(pos_instances > 1.0) assert torch.all(pos_instances < 5.0) + def test_bag_class_probability(): """Verify that pos_class_prob controls the class balance.""" torch.manual_seed(10) - + ds_all_pos = MCStandardMILDataset(D=2, num_bags=10, pos_class_prob=1.0) labels_all_pos = [ds_all_pos[i]["Y"].item() for i in range(10)] assert all(label == 1 for label in labels_all_pos) ds_all_neg = MCStandardMILDataset(D=2, num_bags=10, pos_class_prob=0.0) labels_all_neg = [ds_all_neg[i]["Y"].item() for i in range(10)] - assert all(label == 0 for label in labels_all_neg) \ No newline at end of file + assert all(label == 0 for label in labels_all_neg) diff --git a/tests/datasets/test_trident_wsi_dataset.py b/tests/datasets/test_trident_wsi_dataset.py index ee0fb7b..21e11a8 100644 --- a/tests/datasets/test_trident_wsi_dataset.py +++ b/tests/datasets/test_trident_wsi_dataset.py @@ -7,6 +7,7 @@ # --- Fixtures --- + @pytest.fixture def base_kwargs(): """Standard arguments for initializing the dataset.""" @@ -17,11 +18,13 @@ def base_kwargs(): "magnification": 20, "patch_size": 512, "overlap_pixels": 0, - "wsi_names": ["slide_1", "slide_2"] + "wsi_names": ["slide_1", "slide_2"], } + # --- Tests --- + @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) def test_init_defaults(mock_super_init, base_kwargs): ds = TridentWSIDataset(**base_kwargs) @@ -30,92 +33,106 @@ def test_init_defaults(mock_super_init, base_kwargs): assert ds.magnification == 20 assert ds.overlap_pixels == 0 assert ds.trident_folder == "20x_512px_0px_overlap/" - + expected_features_path = "/mock/base20x_512px_0px_overlap/features_conch_v15/" expected_coords_path = "/mock/base20x_512px_0px_overlap/patches/" - + mock_super_init.assert_called_once() call_kwargs = mock_super_init.call_args[1] - + assert call_kwargs["features_path"] == expected_features_path assert call_kwargs["coords_path"] == expected_coords_path assert call_kwargs["dist_thr"] == pytest.approx(np.sqrt(2.0)) + @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) def test_init_custom_threshold(mock_super_init, base_kwargs): kwargs = base_kwargs.copy() kwargs["dist_thr"] = 5.5 - + # Fix: Assign to '_' to silence the "unused variable" error # while still ensuring TridentWSIDataset initializes. _ = TridentWSIDataset(**kwargs) - + call_kwargs = mock_super_init.call_args[1] assert call_kwargs["dist_thr"] == 5.5 + @patch("torchmil.datasets.wsi_dataset.WSIDataset._load_labels") @patch("os.path.isdir") @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) -def test_load_labels_directory_mode(mock_super_init, mock_isdir, mock_super_load, base_kwargs): +def test_load_labels_directory_mode( + mock_super_init, mock_isdir, mock_super_load, base_kwargs +): """Test directory mode logic. Manually set labels_path since super().__init__ is mocked.""" mock_isdir.return_value = True - + ds = TridentWSIDataset(**base_kwargs) - ds.labels_path = base_kwargs["labels_path"] # Manually set attribute usually set by super() - + ds.labels_path = base_kwargs[ + "labels_path" + ] # Manually set attribute usually set by super() + expected_label = np.array([0]) mock_super_load.return_value = expected_label - + result = ds._load_labels("slide_1") - + assert result == expected_label mock_isdir.assert_called_with("/mock/labels") mock_super_load.assert_called_with("slide_1") + @patch("pandas.read_csv") @patch("os.path.isdir") @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) -def test_load_labels_csv_mode_success(mock_super_init, mock_isdir, mock_read_csv, base_kwargs): +def test_load_labels_csv_mode_success( + mock_super_init, mock_isdir, mock_read_csv, base_kwargs +): """Test CSV mode logic. Manually set labels_path.""" mock_isdir.return_value = False - - df = pd.DataFrame({ - "filename": ["slide_1", "slide_2"], - "grade": [0, 1] - }) + + df = pd.DataFrame({"filename": ["slide_1", "slide_2"], "grade": [0, 1]}) mock_read_csv.return_value = df - + kwargs = base_kwargs.copy() kwargs["wsi_name_col"] = "filename" kwargs["wsi_label_col"] = "grade" - + ds = TridentWSIDataset(**kwargs) - ds.labels_path = kwargs["labels_path"] # Manually set attribute - + ds.labels_path = kwargs["labels_path"] # Manually set attribute + label = ds._load_labels("slide_2") - + assert label[0] == 1 mock_read_csv.assert_called_once_with("/mock/labels") + @patch("pandas.read_csv") @patch("os.path.isdir") @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) -def test_load_labels_csv_missing_kwargs(mock_super_init, mock_isdir, mock_read_csv, base_kwargs): +def test_load_labels_csv_missing_kwargs( + mock_super_init, mock_isdir, mock_read_csv, base_kwargs +): mock_isdir.return_value = False mock_read_csv.return_value = pd.DataFrame() - + ds = TridentWSIDataset(**base_kwargs) - ds.labels_path = base_kwargs["labels_path"] # Manually set attribute - - with pytest.raises(ValueError, match="must provide 'wsi_name_col' and 'wsi_label_col'"): + ds.labels_path = base_kwargs["labels_path"] # Manually set attribute + + with pytest.raises( + ValueError, match="must provide 'wsi_name_col' and 'wsi_label_col'" + ): ds._load_labels("slide_1") + @patch("pandas.read_csv") @patch("os.path.isdir") @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) -def test_load_labels_csv_not_found(mock_super_init, mock_isdir, mock_read_csv, base_kwargs): +def test_load_labels_csv_not_found( + mock_super_init, mock_isdir, mock_read_csv, base_kwargs +): mock_isdir.return_value = False - + # Mock DF to raise ValueError when accessing specific data mock_df = MagicMock() mock_df.loc.__getitem__.side_effect = ValueError("Forced error") @@ -124,13 +141,14 @@ def test_load_labels_csv_not_found(mock_super_init, mock_isdir, mock_read_csv, b kwargs = base_kwargs.copy() kwargs["wsi_name_col"] = "name" kwargs["wsi_label_col"] = "label" - + ds = TridentWSIDataset(**kwargs) - ds.labels_path = kwargs["labels_path"] # Manually set attribute - + ds.labels_path = kwargs["labels_path"] # Manually set attribute + with pytest.raises(ValueError, match="Could not read the label"): ds._load_labels("slide_X") + @patch("h5py.File") @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) def test_load_coords_calculation(mock_super_init, mock_h5, base_kwargs): @@ -138,23 +156,21 @@ def test_load_coords_calculation(mock_super_init, mock_h5, base_kwargs): ds.coords_path = "/mock/patches/" ds.file_type = ".h5" ds.patch_size = 512 - - raw_coords = np.array([ - [1024, 2048], - [1536, 2560] - ]) - + + raw_coords = np.array([[1024, 2048], [1536, 2560]]) + mock_file = MagicMock() mock_file.__getitem__.return_value = raw_coords mock_h5.return_value = mock_file - + expected_coords = np.array([[0, 0], [1, 1]]) - + result = ds._load_coords("slide_1") - + mock_h5.assert_called_with("/mock/patches/slide_1_patches.h5", "r") np.testing.assert_array_equal(result, expected_coords) + @patch("h5py.File") @patch("torchmil.datasets.wsi_dataset.WSIDataset.__init__", return_value=None) def test_load_coords_none(mock_super_init, mock_h5, base_kwargs): @@ -165,19 +181,19 @@ def test_load_coords_none(mock_super_init, mock_h5, base_kwargs): ds = TridentWSIDataset(**base_kwargs) ds.coords_path = "/mock/patches/" ds.file_type = ".h5" - + # 1. Mock the File object mock_file_obj = MagicMock() - + # 2. Mock the dataset object returned by file['coords'] mock_dataset = MagicMock() - + # 3. Mock the slice operator [:] on the dataset to return None mock_dataset.__getitem__.return_value = None - + # 4. Connect them mock_file_obj.__getitem__.return_value = mock_dataset mock_h5.return_value = mock_file_obj - + result = ds._load_coords("slide_empty") - assert result is None \ No newline at end of file + assert result is None diff --git a/tests/nn/test_variational_autoencoder.py b/tests/nn/test_variational_autoencoder.py index b386fac..6edba97 100644 --- a/tests/nn/test_variational_autoencoder.py +++ b/tests/nn/test_variational_autoencoder.py @@ -8,74 +8,85 @@ # --- Fixtures --- + @pytest.fixture def input_dim(): return 10 + @pytest.fixture def latent_dim(): return 5 + @pytest.fixture def layer_sizes(latent_dim): return [8, latent_dim] + @pytest.fixture def sample_data(input_dim): return torch.randn(4, input_dim) # batch_size=4, input_dim=10 + @pytest.fixture def sample_image_data(): # Batch=2, Channels=3, H=4, W=4 -> Flattened size = 48 return torch.randn(2, 3, 4, 4) + @pytest.fixture def sample_bag_data(input_dim): return torch.randn(2, 3, input_dim) # batch_size=2, bag_size=3, input_dim=10 + @pytest.fixture def vae_basic(input_dim, layer_sizes): return VariationalAutoEncoder( - input_shape=(input_dim,), - layer_sizes=layer_sizes, + input_shape=(input_dim,), + layer_sizes=layer_sizes, activations=["relu", "None"], - covar_mode="single" + covar_mode="single", ) + @pytest.fixture def vae_diagonal(input_dim, layer_sizes): return VariationalAutoEncoder( - input_shape=(input_dim,), - layer_sizes=layer_sizes, - covar_mode="diagonal" + input_shape=(input_dim,), layer_sizes=layer_sizes, covar_mode="diagonal" ) + @pytest.fixture def vae_mil(input_dim, layer_sizes): return VariationalAutoEncoderMIL( - input_shape=(input_dim,), - layer_sizes=layer_sizes, - activations=["relu", "None"] + input_shape=(input_dim,), layer_sizes=layer_sizes, activations=["relu", "None"] ) + # --- Tests for VariationalAutoEncoder (Standard) --- + def test_vae_init_validations(): """Test initialization logic and error handling.""" # Test valid diagonal init - vae = VariationalAutoEncoder(input_shape=(10,), layer_sizes=[5], covar_mode="diagonal") + vae = VariationalAutoEncoder( + input_shape=(10,), layer_sizes=[5], covar_mode="diagonal" + ) assert vae.d_var_enc == 5 # Last layer size - assert vae.d_var_dec == 10 # Input shape - + assert vae.d_var_dec == 10 # Input shape + # Test invalid covar mode with pytest.raises(NotImplementedError, match="not valid"): VariationalAutoEncoder(input_shape=(10,), covar_mode="invalid_mode") + def test_vae_forward_standard(sample_data, vae_basic, latent_dim): """Test standard forward pass and shape.""" samples = vae_basic(sample_data, n_samples=2) assert samples.shape == (4, 2, latent_dim) + def test_vae_forward_image_input(sample_image_data): """ Test the flattening logic in forward/get_posterior_samples. @@ -83,21 +94,27 @@ def test_vae_forward_image_input(sample_image_data): """ flat_dim = 3 * 4 * 4 vae = VariationalAutoEncoder(input_shape=(flat_dim,), layer_sizes=[10]) - + # This hits the `if len(X.shape) > 3` block samples = vae(sample_image_data, n_samples=1) assert samples.shape == (2, 1, 10) + def test_vae_forward_returns_stats(sample_data, vae_basic, latent_dim): """Test forward with return_mean_logstd=True.""" - samples, mean, log_std = vae_basic(sample_data, n_samples=1, return_mean_logstd=True) + samples, mean, log_std = vae_basic( + sample_data, n_samples=1, return_mean_logstd=True + ) assert samples.shape == (4, 1, latent_dim) assert mean.shape == (4, latent_dim) # In 'single' mode, log_std returned by forward is expanded to match mean shape logic? # Looking at code: `log_std_v = torch.ones_like(mean) * log_std` - assert log_std.shape == (4, latent_dim) + assert log_std.shape == (4, latent_dim) + -def test_vae_diagonal_covariance_logic(sample_data, vae_diagonal, input_dim, latent_dim): +def test_vae_diagonal_covariance_logic( + sample_data, vae_diagonal, input_dim, latent_dim +): """ Test flow specifically for diagonal covariance mode. Verifies dimensions of variances in encoder and decoder. @@ -105,13 +122,14 @@ def test_vae_diagonal_covariance_logic(sample_data, vae_diagonal, input_dim, lat # 1. Encoder Raw Output mean, log_std = vae_diagonal.get_raw_output_enc(sample_data) assert mean.shape == (4, latent_dim) - assert log_std.shape == (4, latent_dim) # Diagonal mode: var dim == latent dim - + assert log_std.shape == (4, latent_dim) # Diagonal mode: var dim == latent dim + # 2. Decoder Raw Output latent_sample = torch.randn(4, latent_dim) dec_mean, dec_log_std = vae_diagonal.get_raw_output_dec(latent_sample) assert dec_mean.shape == (4, input_dim) - assert dec_log_std.shape == (4, input_dim) # Diagonal mode: var dim == input dim + assert dec_log_std.shape == (4, input_dim) # Diagonal mode: var dim == input dim + def test_vae_complete_forward_samples(sample_data, vae_basic): """Test reconstruction path.""" @@ -119,21 +137,23 @@ def test_vae_complete_forward_samples(sample_data, vae_basic): # Result is averaged over samples assert recs.shape == sample_data.shape + def test_vae_compute_loss_variants(sample_data, vae_basic): """Test all reduction modes and return flags in compute_loss.""" # 1. Reduction = Sum loss_sum = vae_basic.compute_loss(sample_data, reduction="sum") assert loss_sum["VaeELL"].ndim == 0 - + # 2. Reduction = None (returns per instance) loss_none = vae_basic.compute_loss(sample_data, reduction="none") assert loss_none["VaeELL"].shape == (4,) assert loss_none["VaeKL"].shape == (4,) - + # 3. Return Samples loss_dict, samples = vae_basic.compute_loss(sample_data, return_samples=True) assert "VaeELL" in loss_dict - assert samples.shape[0] == 4 * 1 # Batch * n_samples + assert samples.shape[0] == 4 * 1 # Batch * n_samples + def test_vae_compute_loss_image_flattening(sample_image_data): """Test that compute_loss handles >2D input (flattening).""" @@ -143,79 +163,93 @@ def test_vae_compute_loss_image_flattening(sample_image_data): loss = vae.compute_loss(sample_image_data) assert "VaeELL" in loss + def test_vae_importance_sampling(sample_data, vae_basic): """Test log_marginal_likelihood_importance_sampling.""" - log_imp = vae_basic.log_marginal_likelihood_importance_sampling(sample_data, n_samples=10) + log_imp = vae_basic.log_marginal_likelihood_importance_sampling( + sample_data, n_samples=10 + ) assert log_imp.shape == (4,) - + # Test with image data (flattening check) flat_dim = 3 * 4 * 4 img_data = torch.randn(2, 3, 4, 4) vae_img = VariationalAutoEncoder(input_shape=(flat_dim,), layer_sizes=[10]) - log_imp_img = vae_img.log_marginal_likelihood_importance_sampling(img_data, n_samples=2) + log_imp_img = vae_img.log_marginal_likelihood_importance_sampling( + img_data, n_samples=2 + ) assert log_imp_img.shape == (2,) + # --- Tests for VariationalAutoEncoderMIL (MIL Extension) --- + def test_mil_forward_structure(sample_bag_data, vae_mil, latent_dim): """Test basic MIL forward pass dimensions.""" # Input: (2, 3, 10) -> Output: (2, 3, n_samples, latent) samples = vae_mil(sample_bag_data, n_samples=2) assert samples.shape == (2, 3, 2, latent_dim) + def test_mil_forward_single_instance_edge_case(input_dim, vae_mil): """Test forward pass when input is (BagSize, Dim) instead of (Batch, Bag, Dim).""" - single_bag = torch.randn(5, input_dim) + single_bag = torch.randn(5, input_dim) # The code `if len(X.shape) == 2: X = X.unsqueeze(0)` handles this samples = vae_mil(single_bag, n_samples=1) # Output should be (1, 5, 1, latent) assert samples.shape == (1, 5, 1, vae_mil.layer_sizes[-1]) + def test_mil_forward_return_stats(sample_bag_data, vae_mil, latent_dim): """Test return_mean_logstd in MIL context.""" - samples, mean, log_std = vae_mil(sample_bag_data, n_samples=1, return_mean_logstd=True) + samples, mean, log_std = vae_mil( + sample_bag_data, n_samples=1, return_mean_logstd=True + ) assert mean.shape == (2, 3, latent_dim) assert log_std.shape == (2, 3, latent_dim) + def test_mil_complete_forward(sample_bag_data, vae_mil): """Test reconstruction in MIL context.""" recs = vae_mil.complete_forward_samples(sample_bag_data) assert recs.shape == sample_bag_data.shape + def test_mil_compute_loss_masking(sample_bag_data, vae_mil): """Test loss computation with and without masks, and different reductions.""" # Mask: 1 for valid, 0 for padding. Let's mask the last instance of bag 0. mask = torch.ones(2, 3) - mask[0, 2] = 0 - + mask[0, 2] = 0 + # 1. Reduction Mean loss_mean = vae_mil.compute_loss(sample_bag_data, mask=mask, reduction="mean") assert isinstance(loss_mean["VaeELL"], torch.Tensor) - + # 2. Reduction Sum loss_sum = vae_mil.compute_loss(sample_bag_data, mask=mask, reduction="sum") assert isinstance(loss_sum["VaeELL"], torch.Tensor) - + # 3. Reduction None (should return grid) loss_none = vae_mil.compute_loss(sample_bag_data, mask=mask, reduction="none") assert loss_none["VaeELL"].shape == (2, 3) - + # 4. Return Samples loss_dict, samples = vae_mil.compute_loss(sample_bag_data, return_samples=True) - # Expected sample shape: (Batch, n_samples, BagSize, Latent) + # Expected sample shape: (Batch, n_samples, BagSize, Latent) # Note: Code returns `samples.view(B, n_samples, N, -1)` assert samples.shape == (2, 1, 3, vae_mil.layer_sizes[-1]) + def test_mil_importance_sampling(sample_bag_data, vae_mil): """Test MIL importance sampling with mask.""" mask = torch.ones(2, 3) log_imp = vae_mil.log_marginal_likelihood_importance_sampling( sample_bag_data, mask=mask, n_samples=5 ) - assert log_imp.shape == (2, 3) # Returns (Batch, BagSize) - + assert log_imp.shape == (2, 3) # Returns (Batch, BagSize) + # Test without mask (defaults to ones) log_imp_nomask = vae_mil.log_marginal_likelihood_importance_sampling( sample_bag_data, n_samples=5 ) - assert log_imp_nomask.shape == (2, 3) \ No newline at end of file + assert log_imp_nomask.shape == (2, 3) From 167380361ac2d4bda17e9d8486bf45b3fd9b7927 Mon Sep 17 00:00:00 2001 From: Franblueee Date: Fri, 6 Feb 2026 20:41:47 +0100 Subject: [PATCH 10/10] add datasets documentation --- docs/api/datasets/tadmil_dataset.md | 7 +++++++ docs/api/datasets/video_classification_dataset.md | 7 +++++++ mkdocs.yml | 2 ++ torchmil/datasets/tadmil_dataset.py | 2 +- 4 files changed, 17 insertions(+), 1 deletion(-) create mode 100644 docs/api/datasets/tadmil_dataset.md create mode 100644 docs/api/datasets/video_classification_dataset.md diff --git a/docs/api/datasets/tadmil_dataset.md b/docs/api/datasets/tadmil_dataset.md new file mode 100644 index 0000000..119d0a6 --- /dev/null +++ b/docs/api/datasets/tadmil_dataset.md @@ -0,0 +1,7 @@ +# Traffic Anomaly Detection (TAD) MIL Dataset + +::: torchmil.datasets.TADMILDataset + options: + members: + - __init__ + - __getitem__ \ No newline at end of file diff --git a/docs/api/datasets/video_classification_dataset.md b/docs/api/datasets/video_classification_dataset.md new file mode 100644 index 0000000..6421a07 --- /dev/null +++ b/docs/api/datasets/video_classification_dataset.md @@ -0,0 +1,7 @@ +# Video Classification Dataset + +::: torchmil.datasets.VideoClassificationDataset + options: + members: + - __init__ + - __getitem__ diff --git a/mkdocs.yml b/mkdocs.yml index 3fb5491..fd60ac4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -135,9 +135,11 @@ nav: - api/datasets/wsi_dataset.md - api/datasets/trident_wsi_dataset.md - api/datasets/binary_classification_dataset.md + - api/datasets/video_classification_dataset.md - api/datasets/camelyon16mil_dataset.md - api/datasets/pandamil_dataset.md - api/datasets/rsnamil_dataset.md + - api/datasets/tadmil_dataset.md - api/datasets/false_frequency_dataset.md - api/datasets/mc_standard_dataset.md - api/datasets/sc_standard_dataset.md diff --git a/torchmil/datasets/tadmil_dataset.py b/torchmil/datasets/tadmil_dataset.py index 3b570db..16f0c86 100644 --- a/torchmil/datasets/tadmil_dataset.py +++ b/torchmil/datasets/tadmil_dataset.py @@ -15,7 +15,7 @@ class TADMILDataset(BinaryClassificationDataset, VideoClassificationDataset): **Dataset description.** We have preprocessed the Video by computing features for each frame using various feature extractors. - - A **video** is labeled as positive (`frame_label=1`) if it contains evidence of traffic anomaly. + - A **video** is labeled as positive (`label=1`) if it contains evidence of traffic anomaly. - A **video** is labeled as positive (`label=1`) if it contains at least one positive frame. This means a video is considered positive if there is any evidence of traffic anomaly.