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

quadra/datamodules/generic/imagenette.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import os
+import shutil
+from typing import Any
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from torchvision.datasets.utils import download_and_extract_archive
+
+from quadra.datamodules import ClassificationDataModule, SSLDataModule
+from quadra.utils.utils import get_logger
+
+IMAGENETTE_LABEL_MAPPER = {
+    "n01440764": "tench",
+    "n02102040": "english_springer",
+    "n02979186": "cassette_player",
+    "n03000684": "chain_saw",
+    "n03028079": "church",
+    "n03394916": "french_horn",
+    "n03417042": "garbage_truck",
+    "n03425413": "gas_pump",
+    "n03445777": "golf_ball",
+    "n03888257": "parachute",
+}
+
+DEFAULT_CLASS_TO_IDX = {cl: idx for idx, cl in enumerate(sorted(IMAGENETTE_LABEL_MAPPER.values()))}
+
+log = get_logger(__name__)
+
+
+class ImagenetteClassificationDataModule(ClassificationDataModule):
+    """Initializes the classification data module for Imagenette dataset.
+
+    Args:
+        data_path: Path to the dataset.
+        name: Name of the dataset.
+        imagenette_version: Version of the Imagenette dataset. Can be 320 or 160 or full.
+        force_download: If True, the dataset will be downloaded even if the data_path already exists. The data_path
+            will be deleted and recreated.
+        class_to_idx: Dictionary mapping class names to class indices.
+        **kwargs: Keyword arguments for the ClassificationDataModule.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        name: str = "imagenette_classification_datamodule",
+        imagenette_version: str = "320",
+        force_download: bool = False,
+        class_to_idx: dict[str, int] | None = None,
+        **kwargs: Any,
+    ):
+        if imagenette_version not in ["320", "160", "full"]:
+            raise ValueError(f"imagenette_version must be one of 320, 160 or full. Got {imagenette_version} instead.")
+
+        if imagenette_version == "full":
+            imagenette_version = ""
+        else:
+            imagenette_version = f"-{imagenette_version}"
+
+        self.download_url = f"https://s3.amazonaws.com/fast-ai-imageclas/imagenette2{imagenette_version}.tgz"
+        self.force_download = force_download
+        self.imagenette_version = imagenette_version
+
+        if class_to_idx is None:
+            class_to_idx = DEFAULT_CLASS_TO_IDX
+
+        super().__init__(
+            data_path=data_path,
+            name=name,
+            test_split_file=None,
+            train_split_file=None,
+            val_size=None,
+            class_to_idx=class_to_idx,
+            **kwargs,
+        )
+
+    def download_data(self, download_url: str, force_download: bool = False) -> None:
+        """Download the Imagenette dataset.
+
+        Args:
+            download_url: Dataset download url.
+            force_download: If True, the dataset will be downloaded even if the data_path already exists. The data_path
+                will be removed.
+        """
+        if os.path.exists(self.data_path):
+            if force_download:
+                log.info("The path %s already exists. Removing it and downloading the dataset again.", self.data_path)
+                shutil.rmtree(self.data_path)
+            else:
+                log.info("The path %s already exists. Skipping download.", self.data_path)
+                return
+
+        log.info("Downloading and extracting Imagenette dataset to %s", self.data_path)
+        download_and_extract_archive(download_url, self.data_path, remove_finished=True)
+
+    def _prepare_data(self) -> None:
+        """Prepares the data for the data module."""
+        self.download_data(download_url=self.download_url, force_download=self.force_download)
+        self.data_path = os.path.join(self.data_path, f"imagenette2{self.imagenette_version}")
+
+        train_images_and_targets, class_to_idx = self._find_images_and_targets(os.path.join(self.data_path, "train"))
+        self.class_to_idx = {IMAGENETTE_LABEL_MAPPER[k]: v for k, v in class_to_idx.items()}
+
+        samples_train, targets_train = [], []
+        idx_to_class = {v: k for k, v in self.class_to_idx.items()}
+        for image, target in train_images_and_targets:
+            samples_train.append(image)
+            targets_train.append(idx_to_class[target])
+
+        samples_train, samples_val, targets_train, targets_val = train_test_split(
+            samples_train,
+            targets_train,
+            test_size=self.val_size,
+            random_state=self.seed,
+            stratify=targets_train,
+        )
+
+        test_images_and_targets, _ = self._find_images_and_targets(os.path.join(self.data_path, "val"))
+        samples_test, targets_test = [], []
+        for image, target in test_images_and_targets:
+            samples_test.append(image)
+            targets_test.append(idx_to_class[target])
+
+        train_df = pd.DataFrame({"samples": samples_train, "targets": targets_train})
+        train_df["split"] = "train"
+        val_df = pd.DataFrame({"samples": samples_val, "targets": targets_val})
+        val_df["split"] = "val"
+        test_df = pd.DataFrame({"samples": samples_test, "targets": targets_test})
+        test_df["split"] = "test"
+        self.data = pd.concat([train_df, val_df, test_df], axis=0)
+
+
+class ImagenetteSSLDataModule(ImagenetteClassificationDataModule, SSLDataModule):
+    """Initializes the SSL data module for Imagenette dataset."""
+
+    def __init__(
+        self,
+        *args: Any,
+        name="imagenette_ssl",
+        **kwargs: Any,
+    ):
+        super().__init__(*args, name=name, **kwargs)  # type: ignore[misc]

quadra/datamodules/generic/mnist.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import os
+import shutil
+from typing import Any
+
+import cv2
+from torchvision.datasets.mnist import MNIST
+
+from quadra.datamodules import AnomalyDataModule
+from quadra.utils.utils import get_logger
+
+log = get_logger(__name__)
+
+
+class MNISTAnomalyDataModule(AnomalyDataModule):
+    """Standard anomaly datamodule with automatic download of the MNIST dataset."""
+
+    def __init__(
+        self, data_path: str, good_number: int, limit_data: int = 100, category: str | None = None, **kwargs: Any
+    ):
+        """Initialize the MNIST anomaly datamodule.
+
+        Args:
+            data_path: Path to the dataset
+            good_number: Which number to use as a good class, all other numbers are considered anomalies.
+            category: The category of the dataset. For mnist this is always None.
+            limit_data: Limit the number of images to use for training and testing. Defaults to 100.
+            **kwargs: Additional arguments to pass to the AnomalyDataModule.
+        """
+        super().__init__(data_path=data_path, category=None, **kwargs)
+        self.good_number = good_number
+        self.limit_data = limit_data
+
+    def download_data(self) -> None:
+        """Download the MNIST dataset and move images in the right folders."""
+        log.info("Generating MNIST anomaly dataset for good number %s", self.good_number)
+
+        mnist_train_dataset = MNIST(root=self.data_path, train=True, download=True)
+        mnist_test_dataset = MNIST(root=self.data_path, train=False, download=True)
+
+        self.data_path = os.path.join(self.data_path, "quadra_mnist_anomaly")
+
+        if os.path.exists(self.data_path):
+            shutil.rmtree(self.data_path)
+
+        # Create the folder structure
+        train_good_folder = os.path.join(self.data_path, "train", "good")
+        test_good_folder = os.path.join(self.data_path, "test", "good")
+
+        os.makedirs(train_good_folder, exist_ok=True)
+        os.makedirs(test_good_folder, exist_ok=True)
+
+        # Copy the good train images to the correct folder
+        good_train_samples = mnist_train_dataset.data[mnist_train_dataset.targets == self.good_number]
+        for i, image in enumerate(good_train_samples.numpy()):
+            if i == self.limit_data:
+                break
+            cv2.imwrite(os.path.join(train_good_folder, f"{i}.png"), image)
+
+        for number in range(10):
+            if number == self.good_number:
+                good_train_samples = mnist_test_dataset.data[mnist_test_dataset.targets == number]
+                for i, image in enumerate(good_train_samples.numpy()):
+                    if i == self.limit_data:
+                        break
+                    cv2.imwrite(os.path.join(test_good_folder, f"{number}_{i}.png"), image)
+            else:
+                test_bad_folder = os.path.join(self.data_path, "test", str(number))
+                os.makedirs(test_bad_folder, exist_ok=True)
+                bad_train_samples = mnist_train_dataset.data[mnist_train_dataset.targets == number]
+                for i, image in enumerate(bad_train_samples.numpy()):
+                    if i == self.limit_data:
+                        break
+
+                    cv2.imwrite(os.path.join(test_bad_folder, f"{number}_{i}.png"), image)
+
+    def _prepare_data(self) -> None:
+        """Prepare the MNIST dataset."""
+        self.download_data()
+        return super()._prepare_data()

quadra/datamodules/generic/mvtec.py

@@ -0,0 +1,58 @@
+import os
+from pathlib import Path
+
+from torchvision.datasets.utils import download_and_extract_archive
+
+from quadra.datamodules import AnomalyDataModule
+from quadra.utils.utils import get_logger
+
+log = get_logger(__name__)
+
+
+DATASET_BASE_URL = "https://www.mydrive.ch/shares/38536/3830184030e49fe74747669442f0f282/download/"
+
+DATASET_URL = {
+    "bottle": DATASET_BASE_URL + "420937370-1629951468/bottle.tar.xz",
+    "capsule": DATASET_BASE_URL + "420937454-1629951595/capsule.tar.xz",
+    "carpet": DATASET_BASE_URL + "420937484-1629951672/carpet.tar.xz",
+    "grid": DATASET_BASE_URL + "420937487-1629951814/grid.tar.xz",
+    "hazelnut": DATASET_BASE_URL + "420937545-1629951845/hazelnut.tar.xz",
+    "leather": DATASET_BASE_URL + "420937607-1629951964/leather.tar.xz",
+    "metal_nut": DATASET_BASE_URL + "420937637-1629952063/metal_nut.tar.xz",
+    "pill": DATASET_BASE_URL + "420938129-1629953099/pill.tar.xz",
+    "screw": DATASET_BASE_URL + "420938130-1629953152/screw.tar.xz",
+    "tile": DATASET_BASE_URL + "420938133-1629953189/tile.tar.xz",
+    "toothbrush": DATASET_BASE_URL + "420938134-1629953256/toothbrush.tar.xz",
+    "transistor": DATASET_BASE_URL + "420938166-1629953277/transistor.tar.xz",
+    "wood": DATASET_BASE_URL + "420938383-1629953354/wood.tar.xz",
+    "zipper": DATASET_BASE_URL + "420938385-1629953449/zipper.tar.xz",
+}
+
+
+class MVTecDataModule(AnomalyDataModule):
+    """Standard anomaly datamodule with automatic download of the MVTec dataset."""
+
+    def __init__(self, data_path: str, category: str, **kwargs):
+        if category not in DATASET_URL:
+            raise ValueError(f"Unknown category {category}. Available categories are {list(DATASET_URL.keys())}")
+
+        super().__init__(data_path=data_path, category=category, **kwargs)
+
+    def download_data(self) -> None:
+        """Download the MVTec dataset."""
+        if self.category is None:
+            raise ValueError("Category must be specified for MVTec dataset.")
+
+        if os.path.exists(self.data_path):
+            log.info("The path %s already exists. Skipping download.", os.path.join(self.data_path, self.category))
+            return
+
+        log.info("Downloading and extracting MVTec dataset for category %s to %s", self.category, self.data_path)
+        # self.data_path is the path to the category folder that will be created by the download_and_extract_archive
+        data_path_no_category = str(Path(self.data_path).parent)
+        download_and_extract_archive(DATASET_URL[self.category], data_path_no_category, remove_finished=True)
+
+    def _prepare_data(self) -> None:
+        """Prepare the MVTec dataset."""
+        self.download_data()
+        return super()._prepare_data()

quadra/datamodules/generic/oxford_pet.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import albumentations
+import cv2
+import numpy as np
+import pandas as pd
+from torchvision.datasets.utils import download_and_extract_archive
+
+from quadra.datamodules import SegmentationMulticlassDataModule
+from quadra.datasets.segmentation import SegmentationDatasetMulticlass
+from quadra.utils import utils
+
+log = utils.get_logger(__name__)
+
+
+class OxfordPetSegmentationDataModule(SegmentationMulticlassDataModule):
+    """OxfordPetSegmentationDataModule.
+
+    Args:
+        data_path: path to the oxford pet dataset
+        idx_to_class: dict with corrispondence btw mask index and classes: {1: class_1, 2: class_2, ..., N: class_N}
+            except background class which is 0.
+        name: Defaults to "oxford_pet_segmentation_datamodule".
+        dataset: Defaults to SegmentationDataset.
+        batch_size:  batch size for training. Defaults to 32.
+        test_size: Defaults to 0.3.
+        val_size:  Defaults to 0.3.
+        seed: Defaults to 42.
+        num_workers: number of workers for data loading. Defaults to 6.
+        train_transform: Train transform. Defaults to None.
+        test_transform: Test transform. Defaults to None.
+        val_transform: Validation transform. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        idx_to_class: dict,
+        name: str = "oxford_pet_segmentation_datamodule",
+        dataset: type[SegmentationDatasetMulticlass] = SegmentationDatasetMulticlass,
+        batch_size: int = 32,
+        test_size: float = 0.3,
+        val_size: float = 0.3,
+        seed: int = 42,
+        num_workers: int = 6,
+        train_transform: albumentations.Compose | None = None,
+        test_transform: albumentations.Compose | None = None,
+        val_transform: albumentations.Compose | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            data_path=data_path,
+            idx_to_class=idx_to_class,
+            name=name,
+            dataset=dataset,
+            batch_size=batch_size,
+            test_size=test_size,
+            val_size=val_size,
+            seed=seed,
+            num_workers=num_workers,
+            train_transform=train_transform,
+            test_transform=test_transform,
+            val_transform=val_transform,
+            **kwargs,
+        )
+
+    _RESOURCES = (
+        ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz", "5c4f3ee8e5d25df40f4fd59a7f44e54c"),
+        ("https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz", "95a8c909bbe2e81eed6a22bccdf3f68f"),
+    )
+
+    def _preprocess_mask(self, mask: np.ndarray) -> np.ndarray:
+        """Preprocess mask function that is adapted from
+        https://albumentations.ai/docs/examples/pytorch_semantic_segmentation/.
+
+        Args:
+            mask: mask to be preprocessed
+
+        Returns:
+            binarized mask
+        """
+        mask = mask.astype(np.float32)
+        mask[mask == 2.0] = 0.0
+        mask[(mask == 1.0) | (mask == 3.0)] = 1.0
+        mask = (mask > 0).astype(np.uint8)
+        return mask
+
+    def _check_exists(self, image_folder: str, annotation_folder: str) -> bool:
+        """Check if the dataset is already downloaded."""
+        return all(os.path.exists(folder) and os.path.isdir(folder) for folder in (image_folder, annotation_folder))
+
+    def download_data(self):
+        """Download the dataset if it is not already downloaded."""
+        image_folder = os.path.join(self.data_path, "images")
+        annotation_folder = os.path.join(self.data_path, "annotations")
+        if not self._check_exists(image_folder, annotation_folder):
+            for url, md5 in self._RESOURCES:
+                download_and_extract_archive(url, download_root=self.data_path, md5=md5, remove_finished=True)
+            log.info("Fixing corrupted files...")
+            images_filenames = sorted(os.listdir(image_folder))
+            for filename in images_filenames:
+                file_wo_ext = os.path.splitext(os.path.basename(filename))[0]
+                try:
+                    mask = cv2.imread(os.path.join(annotation_folder, "trimaps", file_wo_ext + ".png"))
+                    mask = self._preprocess_mask(mask)
+                    if np.sum(mask) == 0:
+                        os.remove(os.path.join(image_folder, filename))
+                        os.remove(os.path.join(annotation_folder, "trimaps", file_wo_ext + ".png"))
+                        log.info("Removed %s", filename)
+                    else:
+                        img = cv2.imread(os.path.join(image_folder, filename))
+                        cv2.imwrite(os.path.join(image_folder, file_wo_ext + ".jpg"), img)
+                except Exception:
+                    ip = os.path.join(image_folder, filename)
+                    mp = os.path.join(annotation_folder, "trimaps", file_wo_ext + ".png")
+                    if os.path.exists(ip):
+                        os.remove(ip)
+                    if os.path.exists(mp):
+                        os.remove(mp)
+                    log.info("Removed %s", filename)
+
+    def _prepare_data(self) -> None:
+        """Prepare the data to be used by the DataModule."""
+        self.download_data()
+
+        trainval_split_filepath = os.path.join(self.data_path, "annotations", "trainval.txt")
+        with open(trainval_split_filepath) as f:
+            split_data = f.read().strip("\n").split("\n")
+        trainval_filenames = [
+            x.split(" ")[0]
+            for x in split_data
+            if os.path.exists(os.path.join(self.data_path, "images", x.split(" ")[0] + ".jpg"))
+        ]
+        train_filenames = [x for i, x in enumerate(trainval_filenames) if i % 10 != 0]
+        val_filenames = [x for i, x in enumerate(trainval_filenames) if i % 10 == 0]
+
+        test_split_filepath = os.path.join(self.data_path, "annotations", "test.txt")
+        with open(test_split_filepath) as f:
+            split_data = f.read().strip("\n").split("\n")
+        test_filenames = [
+            x.split(" ")[0]
+            for x in split_data
+            if os.path.exists(os.path.join(self.data_path, "images", x.split(" ")[0] + ".jpg"))
+        ]
+
+        df_list = []
+        for split_name, filenames in [
+            ("train", train_filenames),
+            ("val", val_filenames),
+            ("test", test_filenames),
+        ]:
+            samples = [os.path.join(self.data_path, "images", f + ".jpg") for f in filenames]
+            masks = [os.path.join(self.data_path, "annotations", "trimaps", f + ".png") for f in filenames]
+            targets = [1] * len(filenames)
+
+            df = pd.DataFrame({"samples": samples, "masks": masks, "targets": targets})
+            df["split"] = split_name
+            df_list.append(df)
+
+        self.data = pd.concat(df_list, axis=0)

quadra/datamodules/patch.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import json
+import os
+
+import albumentations
+import pandas as pd
+from torch.utils.data import DataLoader
+
+from quadra.datamodules.base import BaseDataModule
+from quadra.datasets import ImageClassificationListDataset, PatchSklearnClassificationTrainDataset
+from quadra.utils.classification import find_test_image
+from quadra.utils.patch.dataset import PatchDatasetInfo, load_train_file
+
+
+class PatchSklearnClassificationDataModule(BaseDataModule):
+    """DataModule for patch classification.
+
+    Args:
+        data_path: Location of the dataset
+        name: Name of the datamodule
+        train_filename: Name of the file containing the list of training samples
+        exclude_filter: Filter to exclude samples from the dataset
+        include_filter: Filter to include samples from the dataset
+        class_to_idx: Dictionary mapping class names to indices
+        seed: Random seed
+        batch_size: Batch size
+        num_workers: Number of workers
+        train_transform: Transform to apply to the training samples
+        val_transform: Transform to apply to the validation samples
+        test_transform: Transform to apply to the test samples
+        balance_classes: If True repeat low represented classes
+        class_to_skip_training: List of classes skipped during training.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        class_to_idx: dict,
+        name: str = "patch_classification_datamodule",
+        train_filename: str = "dataset.txt",
+        exclude_filter: list[str] | None = None,
+        include_filter: list[str] | None = None,
+        seed: int = 42,
+        batch_size: int = 32,
+        num_workers: int = 6,
+        train_transform: albumentations.Compose | None = None,
+        val_transform: albumentations.Compose | None = None,
+        test_transform: albumentations.Compose | None = None,
+        balance_classes: bool = False,
+        class_to_skip_training: list | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            data_path=data_path,
+            name=name,
+            seed=seed,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            train_transform=train_transform,
+            val_transform=val_transform,
+            test_transform=test_transform,
+            **kwargs,
+        )
+        self.class_to_idx = class_to_idx
+        self.balance_classes = balance_classes
+        self.train_filename = train_filename
+        self.include_filter = include_filter
+        self.exclude_filter = exclude_filter
+        self.class_to_skip_training = class_to_skip_training
+
+        self.train_folder = os.path.join(self.data_path, "train")
+        self.val_folder = os.path.join(self.data_path, "val")
+        self.test_folder = os.path.join(self.data_path, "test")
+        self.info: PatchDatasetInfo
+        self.train_dataset: PatchSklearnClassificationTrainDataset
+        self.val_dataset: ImageClassificationListDataset
+        self.test_dataset: ImageClassificationListDataset
+
+    def _prepare_data(self):
+        """Prepare data function."""
+        if os.path.isfile(os.path.join(self.data_path, "info.json")):
+            with open(os.path.join(self.data_path, "info.json")) as f:
+                self.info = PatchDatasetInfo(**json.load(f))
+        else:
+            raise FileNotFoundError("No `info.json` file found in the dataset folder")
+
+        split_df_list: list[pd.DataFrame] = []
+        if os.path.isfile(os.path.join(self.train_folder, self.train_filename)):
+            train_samples, train_labels = load_train_file(
+                train_file_path=os.path.join(self.train_folder, self.train_filename),
+                include_filter=self.include_filter,
+                exclude_filter=self.exclude_filter,
+                class_to_skip=self.class_to_skip_training,
+            )
+            train_df = pd.DataFrame({"samples": train_samples, "targets": train_labels})
+            train_df["split"] = "train"
+            split_df_list.append(train_df)
+        if os.path.isdir(self.val_folder):
+            val_samples, val_labels = find_test_image(
+                folder=self.val_folder,
+                exclude_filter=self.exclude_filter,
+                include_filter=self.include_filter,
+                include_none_class=False,
+            )
+            val_df = pd.DataFrame({"samples": val_samples, "targets": val_labels})
+            val_df["split"] = "val"
+            split_df_list.append(val_df)
+        if os.path.isdir(self.test_folder):
+            test_samples, test_labels = find_test_image(
+                folder=self.test_folder,
+                exclude_filter=self.exclude_filter,
+                include_filter=self.include_filter,
+                include_none_class=True,
+            )
+            test_df = pd.DataFrame({"samples": test_samples, "targets": test_labels})
+            test_df["split"] = "test"
+            split_df_list.append(test_df)
+        if len(split_df_list) == 0:
+            raise ValueError("No data found in all split folders")
+        self.data = pd.concat(split_df_list, axis=0)
+
+    def setup(self, stage: str | None = None) -> None:
+        """Setup function."""
+        if stage == "fit":
+            self.train_dataset = PatchSklearnClassificationTrainDataset(
+                data_path=self.data_path,
+                class_to_idx=self.class_to_idx,
+                samples=self.data[self.data["split"] == "train"]["samples"].tolist(),
+                targets=self.data[self.data["split"] == "train"]["targets"].tolist(),
+                transform=self.train_transform,
+                balance_classes=self.balance_classes,
+            )
+
+            self.val_dataset = ImageClassificationListDataset(
+                class_to_idx=self.class_to_idx,
+                samples=self.data[self.data["split"] == "val"]["samples"].tolist(),
+                targets=self.data[self.data["split"] == "val"]["targets"].tolist(),
+                transform=self.val_transform,
+                allow_missing_label=False,
+            )
+
+        elif stage in ["test", "predict"]:
+            self.test_dataset = ImageClassificationListDataset(
+                class_to_idx=self.class_to_idx,
+                samples=self.data[self.data["split"] == "test"]["samples"].tolist(),
+                targets=self.data[self.data["split"] == "test"]["targets"].tolist(),
+                transform=self.test_transform,
+                allow_missing_label=True,
+            )
+
+    def train_dataloader(self) -> DataLoader:
+        """Return the train dataloader."""
+        if not self.train_dataset_available:
+            raise ValueError("No training sample is available")
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        """Return the validation dataloader."""
+        if not self.val_dataset_available:
+            raise ValueError("No validation dataset is available")
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+        )
+
+    def test_dataloader(self) -> DataLoader:
+        """Return the test dataloader."""
+        if not self.test_dataset_available:
+            raise ValueError("No test dataset is available")
+
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+        )