diff --git a/configs/data/avenue.yaml b/configs/data/avenue.yaml index 396a9ba6b5..8fb07660ce 100644 --- a/configs/data/avenue.yaml +++ b/configs/data/avenue.yaml @@ -8,9 +8,6 @@ init_args: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - transform: null - train_transform: null - eval_transform: null val_split_mode: from_test val_split_ratio: 0.5 seed: null diff --git a/configs/data/btech.yaml b/configs/data/btech.yaml index 22bfd0d8fe..9aa030540c 100644 --- a/configs/data/btech.yaml +++ b/configs/data/btech.yaml @@ -5,9 +5,6 @@ init_args: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test diff --git a/configs/data/folder.yaml b/configs/data/folder.yaml index 329fba6520..76be1382a7 100644 --- a/configs/data/folder.yaml +++ b/configs/data/folder.yaml @@ -12,9 +12,6 @@ init_args: eval_batch_size: 32 num_workers: 8 task: segmentation - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test diff --git a/configs/data/kolektor.yaml b/configs/data/kolektor.yaml index 1b2e6fe6b4..5daec435e4 100644 --- a/configs/data/kolektor.yaml +++ b/configs/data/kolektor.yaml @@ -4,9 +4,6 @@ init_args: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test diff --git a/configs/data/mvtec.yaml b/configs/data/mvtec.yaml index 7728808ece..5fb206e144 100644 --- a/configs/data/mvtec.yaml +++ b/configs/data/mvtec.yaml @@ -6,9 +6,6 @@ init_args: eval_batch_size: 32 num_workers: 8 task: segmentation - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test diff --git a/configs/data/mvtec_3d.yaml b/configs/data/mvtec_3d.yaml index d880f92f8f..f567f80899 100644 --- a/configs/data/mvtec_3d.yaml +++ b/configs/data/mvtec_3d.yaml @@ -5,9 +5,6 @@ init_args: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test diff --git a/configs/data/shanghaitech.yaml b/configs/data/shanghaitech.yaml index be4da54311..d18e7671dc 100644 --- a/configs/data/shanghaitech.yaml +++ b/configs/data/shanghaitech.yaml @@ -8,9 +8,6 @@ init_args: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - transform: null - train_transform: null - eval_transform: null val_split_mode: FROM_TEST val_split_ratio: 0.5 seed: null diff --git a/configs/data/ucsd_ped.yaml b/configs/data/ucsd_ped.yaml index 009a5ef224..1226e4f149 100644 --- a/configs/data/ucsd_ped.yaml +++ b/configs/data/ucsd_ped.yaml @@ -8,9 +8,6 @@ init_args: train_batch_size: 8 eval_batch_size: 1 num_workers: 8 - transform: null - train_transform: null - eval_transform: null val_split_mode: FROM_TEST val_split_ratio: 0.5 seed: null diff --git a/configs/data/visa.yaml b/configs/data/visa.yaml index c5656a2158..0d94e82fa4 100644 --- a/configs/data/visa.yaml +++ b/configs/data/visa.yaml @@ -5,9 +5,6 @@ init_args: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test diff --git a/notebooks/100_datamodules/101_btech.ipynb b/notebooks/100_datamodules/101_btech.ipynb index ef188665e6..2b87763ff0 100644 --- a/notebooks/100_datamodules/101_btech.ipynb +++ b/notebooks/100_datamodules/101_btech.ipynb @@ -48,7 +48,7 @@ "# NOTE: Provide the path to the dataset root directory.\n", "# If the datasets is not downloaded, it will be downloaded\n", "# to this directory.\n", - "dataset_root = Path.cwd().parent / \"datasets\" / \"BTech\"" + "dataset_root = Path.cwd().parent.parent / \"datasets\" / \"BTech\"" ] }, { @@ -106,7 +106,6 @@ "btech_datamodule = BTech(\n", " root=dataset_root,\n", " category=\"01\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=0,\n", @@ -378,7 +377,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.14" + "version": "3.11.8" }, "orig_nbformat": 4, "vscode": { diff --git a/notebooks/100_datamodules/102_mvtec.ipynb b/notebooks/100_datamodules/102_mvtec.ipynb index 4c274939d6..9081f256ae 100644 --- a/notebooks/100_datamodules/102_mvtec.ipynb +++ b/notebooks/100_datamodules/102_mvtec.ipynb @@ -58,7 +58,7 @@ "# NOTE: Provide the path to the dataset root directory.\n", "# If the datasets is not downloaded, it will be downloaded\n", "# to this directory.\n", - "dataset_root = Path.cwd().parent / \"datasets\" / \"MVTec\"" + "dataset_root = Path.cwd().parent.parent / \"datasets\" / \"MVTec\"" ] }, { @@ -84,7 +84,6 @@ "mvtec_datamodule = MVTec(\n", " root=dataset_root,\n", " category=\"bottle\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=0,\n", @@ -345,7 +344,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.14" + "version": "3.11.8" }, "orig_nbformat": 4, "vscode": { diff --git a/notebooks/100_datamodules/103_folder.ipynb b/notebooks/100_datamodules/103_folder.ipynb index 2f642e145a..328a069652 100644 --- a/notebooks/100_datamodules/103_folder.ipynb +++ b/notebooks/100_datamodules/103_folder.ipynb @@ -33,7 +33,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 28, "metadata": {}, "outputs": [], "source": [ @@ -42,7 +42,7 @@ "# NOTE: Provide the path to the dataset root directory.\n", "# If the datasets is not downloaded, it will be downloaded\n", "# to this directory.\n", - "dataset_root = Path.cwd().parent / \"datasets\" / \"hazelnut_toy\"" + "dataset_root = Path.cwd().parent.parent / \"datasets\" / \"hazelnut_toy\"" ] }, { @@ -63,7 +63,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 29, "metadata": {}, "outputs": [], "source": [ @@ -91,7 +91,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 30, "metadata": {}, "outputs": [], "source": [ @@ -102,7 +102,6 @@ " abnormal_dir=\"crack\",\n", " task=TaskType.SEGMENTATION,\n", " mask_dir=dataset_root / \"mask\" / \"crack\",\n", - " image_size=(256, 256),\n", ")\n", "folder_datamodule.setup()" ] @@ -114,7 +113,7 @@ "outputs": [], "source": [ "# Train images\n", - "i, data = next(enumerate(folder_datamodule.train_dataloader()))\n", + "data = next(iter(folder_datamodule.train_data))\n", "print(data.image.shape)" ] }, @@ -125,7 +124,7 @@ "outputs": [], "source": [ "# Test images\n", - "i, data = next(enumerate(folder_datamodule.test_dataloader()))\n", + "data = next(iter(folder_datamodule.test_data))\n", "print(data.image.shape, data.gt_mask.shape)" ] }, @@ -143,8 +142,8 @@ "metadata": {}, "outputs": [], "source": [ - "img = to_pil_image(data.image[0].clone())\n", - "msk = to_pil_image(data.gt_mask[0].int() * 255).convert(\"RGB\")\n", + "img = to_pil_image(data.image.clone())\n", + "msk = to_pil_image(data.gt_mask.int() * 255).convert(\"RGB\")\n", "\n", "Image.fromarray(np.hstack((np.array(img), np.array(msk))))" ] @@ -187,7 +186,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 36, "metadata": {}, "outputs": [], "source": [ @@ -369,7 +368,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.14" + "version": "3.11.8" }, "orig_nbformat": 4, "vscode": { diff --git a/notebooks/100_datamodules/104_tiling.ipynb b/notebooks/100_datamodules/104_tiling.ipynb index 949d6f1cf1..dd901c37e7 100644 --- a/notebooks/100_datamodules/104_tiling.ipynb +++ b/notebooks/100_datamodules/104_tiling.ipynb @@ -44,7 +44,7 @@ "# NOTE: Provide the path to the dataset root directory.\n", "# If the datasets is not downloaded, it will be downloaded\n", "# to this directory.\n", - "dataset_root = Path.cwd().parent / \"datasets\" / \"MVTec\" / \"transistor\"" + "dataset_root = Path.cwd().parent.parent / \"datasets\" / \"MVTec\" / \"transistor\"" ] }, { diff --git a/notebooks/200_models/201_fastflow.ipynb b/notebooks/200_models/201_fastflow.ipynb index ad93049ac0..492655f010 100644 --- a/notebooks/200_models/201_fastflow.ipynb +++ b/notebooks/200_models/201_fastflow.ipynb @@ -44,7 +44,7 @@ "# NOTE: Provide the path to the dataset root directory.\n", "# If the datasets is not downloaded, it will be downloaded\n", "# to this directory.\n", - "dataset_root = Path.cwd().parent / \"datasets\" / \"MVTec\"" + "dataset_root = Path.cwd().parent.parent / \"datasets\" / \"MVTec\"" ] }, { @@ -120,7 +120,6 @@ "datamodule = MVTec(\n", " root=dataset_root,\n", " category=\"bottle\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=0,\n", @@ -319,7 +318,9 @@ }, "outputs": [], "source": [ - "inference_dataset = PredictDataset(path=dataset_root / \"bottle/test/broken_large/000.png\")\n", + "pre_processor = Fastflow.configure_pre_processor()\n", + "transform = pre_processor.predict_transform\n", + "inference_dataset = PredictDataset(path=dataset_root / \"bottle/test/broken_large/000.png\", transform=transform)\n", "inference_dataloader = DataLoader(dataset=inference_dataset, collate_fn=inference_dataset.collate_fn)" ] }, @@ -554,7 +555,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.14" + "version": "3.11.8" }, "orig_nbformat": 4, "vscode": { diff --git a/notebooks/600_loggers/601_mlflow_logging.ipynb b/notebooks/600_loggers/601_mlflow_logging.ipynb index b6b3c424cc..c3f37de763 100644 --- a/notebooks/600_loggers/601_mlflow_logging.ipynb +++ b/notebooks/600_loggers/601_mlflow_logging.ipynb @@ -135,7 +135,7 @@ "# NOTE: Provide the path to the dataset root directory.\n", "# If the datasets is not downloaded, it will be downloaded\n", "# to this directory.\n", - "dataset_root = Path.cwd().parent / \"datasets\" / \"MVTec\"" + "dataset_root = Path.cwd().parent.parent / \"datasets\" / \"MVTec\"" ] }, { @@ -197,7 +197,6 @@ "datamodule = MVTec(\n", " root=dataset_root,\n", " category=\"bottle\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=24,\n", @@ -420,7 +419,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.14" + "version": "3.11.8" } }, "nbformat": 4, diff --git a/notebooks/700_metrics/701a_aupimo.ipynb b/notebooks/700_metrics/701a_aupimo.ipynb index ecafbbb7ba..c3be846a77 100644 --- a/notebooks/700_metrics/701a_aupimo.ipynb +++ b/notebooks/700_metrics/701a_aupimo.ipynb @@ -140,7 +140,6 @@ "datamodule = MVTec(\n", " root=dataset_root,\n", " category=\"leather\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=8,\n", @@ -405,7 +404,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.14" + "version": "3.11.8" }, "orig_nbformat": 4 }, diff --git a/notebooks/700_metrics/701b_aupimo_advanced_i.ipynb b/notebooks/700_metrics/701b_aupimo_advanced_i.ipynb index 42c45e60aa..70f0968520 100644 --- a/notebooks/700_metrics/701b_aupimo_advanced_i.ipynb +++ b/notebooks/700_metrics/701b_aupimo_advanced_i.ipynb @@ -164,7 +164,6 @@ "datamodule = MVTec(\n", " root=dataset_root,\n", " category=\"leather\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=8,\n", diff --git a/notebooks/700_metrics/701c_aupimo_advanced_ii.ipynb b/notebooks/700_metrics/701c_aupimo_advanced_ii.ipynb index a83abd5ee2..1d64e9ec44 100644 --- a/notebooks/700_metrics/701c_aupimo_advanced_ii.ipynb +++ b/notebooks/700_metrics/701c_aupimo_advanced_ii.ipynb @@ -158,7 +158,6 @@ "datamodule = MVTec(\n", " root=dataset_root,\n", " category=\"leather\",\n", - " image_size=256,\n", " train_batch_size=32,\n", " eval_batch_size=32,\n", " num_workers=8,\n", diff --git a/src/anomalib/data/datamodules/base/image.py b/src/anomalib/data/datamodules/base/image.py index 28fd9499eb..8476bf5eeb 100644 --- a/src/anomalib/data/datamodules/base/image.py +++ b/src/anomalib/data/datamodules/base/image.py @@ -12,7 +12,6 @@ from lightning.pytorch.trainer.states import TrainerFn from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS from torch.utils.data.dataloader import DataLoader -from torchvision.transforms.v2 import Resize, Transform from anomalib.data.utils import TestSplitMode, ValSplitMode, random_split, split_by_label from anomalib.data.utils.synthetic import SyntheticAnomalyDataset @@ -40,14 +39,6 @@ class AnomalibDataModule(LightningDataModule, ABC): Defaults to ``None``. test_split_ratio (float): Fraction of the train images held out for testing. Defaults to ``None``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. seed (int | None, optional): Seed used during random subset splitting. Defaults to ``None``. """ @@ -61,10 +52,6 @@ def __init__( val_split_ratio: float, test_split_mode: TestSplitMode | str | None = None, test_split_ratio: float | None = None, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, seed: int | None = None, ) -> None: super().__init__() @@ -75,18 +62,8 @@ def __init__( self.test_split_ratio = test_split_ratio self.val_split_mode = ValSplitMode(val_split_mode) self.val_split_ratio = val_split_ratio - self.image_size = image_size self.seed = seed - # set transforms - if bool(train_transform) != bool(eval_transform): - msg = "Only one of train_transform and eval_transform was specified. This is not recommended because \ - it could lead to unexpected behaviour. Please ensure training and eval transforms have the same \ - reshape and normalization characteristics." - logger.warning(msg) - self._train_transform = train_transform or transform - self._eval_transform = eval_transform or transform - self.train_data: AnomalibDataset self.val_data: AnomalibDataset self.test_data: AnomalibDataset @@ -228,46 +205,6 @@ def predict_dataloader(self) -> EVAL_DATALOADERS: """Use the test dataloader for inference unless overridden.""" return self.test_dataloader() - @property - def transform(self) -> Transform: - """Property that returns the user-specified transform for the datamodule, if any. - - This property is accessed by the engine to set the transform for the model. The eval_transform takes precedence - over the train_transform, because the transform that we store in the model is the one that should be used during - inference. - """ - if self._eval_transform: - return self._eval_transform - return None - - @property - def train_transform(self) -> Transform: - """Get the transforms that will be passed to the train dataset. - - If the train_transform is not set, the engine will request the transform from the model. - """ - if self._train_transform: - return self._train_transform - if getattr(self, "trainer", None) and self.trainer.lightning_module and self.trainer.lightning_module.transform: - return self.trainer.lightning_module.transform - if self.image_size: - return Resize(self.image_size, antialias=True) - return None - - @property - def eval_transform(self) -> Transform: - """Get the transform that will be passed to the val/test/predict datasets. - - If the eval_transform is not set, the engine will request the transform from the model. - """ - if self._eval_transform: - return self._eval_transform - if getattr(self, "trainer", None) and self.trainer.lightning_module and self.trainer.lightning_module.transform: - return self.trainer.lightning_module.transform - if self.image_size: - return Resize(self.image_size, antialias=True) - return None - @classmethod def from_config( cls: type["AnomalibDataModule"], diff --git a/src/anomalib/data/datamodules/depth/folder_3d.py b/src/anomalib/data/datamodules/depth/folder_3d.py index cebea42d02..2e2930be26 100644 --- a/src/anomalib/data/datamodules/depth/folder_3d.py +++ b/src/anomalib/data/datamodules/depth/folder_3d.py @@ -8,8 +8,6 @@ from pathlib import Path -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.depth.folder_3d import Folder3DDataset @@ -51,14 +49,6 @@ class Folder3D(AnomalibDataModule): Defaults to ``8``. task (TaskType, optional): Task type. Could be ``classification``, ``detection`` or ``segmentation``. Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR``. test_split_ratio (float): Fraction of images from the train set that will be reserved for testing. @@ -87,10 +77,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.FROM_TEST, @@ -101,10 +87,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, @@ -127,7 +109,6 @@ def _setup(self, _stage: str | None = None) -> None: self.train_data = Folder3DDataset( name=self.name, task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.root, normal_dir=self.normal_dir, @@ -143,7 +124,6 @@ def _setup(self, _stage: str | None = None) -> None: self.test_data = Folder3DDataset( name=self.name, task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.root, normal_dir=self.normal_dir, diff --git a/src/anomalib/data/datamodules/depth/mvtec_3d.py b/src/anomalib/data/datamodules/depth/mvtec_3d.py index 1e5b90e917..6a497ec952 100644 --- a/src/anomalib/data/datamodules/depth/mvtec_3d.py +++ b/src/anomalib/data/datamodules/depth/mvtec_3d.py @@ -22,18 +22,10 @@ import logging from pathlib import Path -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.depth.mvtec_3d import MVTec3DDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - TestSplitMode, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, TestSplitMode, ValSplitMode, download_and_extract logger = logging.getLogger(__name__) @@ -62,14 +54,6 @@ class MVTec3D(AnomalibDataModule): Defaults to ``8``. task (TaskType): Task type, 'classification', 'detection' or 'segmentation' Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR``. test_split_ratio (float): Fraction of images from the train set that will be reserved for testing. @@ -90,10 +74,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST, @@ -104,10 +84,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, @@ -122,14 +98,12 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = MVTec3DDataset( task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.root, category=self.category, ) self.test_data = MVTec3DDataset( task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.root, category=self.category, diff --git a/src/anomalib/data/datamodules/image/btech.py b/src/anomalib/data/datamodules/image/btech.py index 5abda6156e..818c9d71b5 100644 --- a/src/anomalib/data/datamodules/image/btech.py +++ b/src/anomalib/data/datamodules/image/btech.py @@ -14,19 +14,12 @@ from pathlib import Path import cv2 -from torchvision.transforms.v2 import Transform from tqdm import tqdm from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.image.btech import BTechDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - TestSplitMode, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, TestSplitMode, ValSplitMode, download_and_extract logger = logging.getLogger(__name__) @@ -53,14 +46,6 @@ class BTech(AnomalibDataModule): Defaults to ``8``. task (TaskType, optional): Task type. Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode, optional): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR``. test_split_ratio (float, optional): Fraction of images from the train set that will be reserved for testing. @@ -79,12 +64,9 @@ class BTech(AnomalibDataModule): >>> datamodule = BTech( ... root="./datasets/BTech", ... category="01", - ... image_size=256, ... train_batch_size=32, ... eval_batch_size=32, ... num_workers=8, - ... transform_config_train=None, - ... transform_config_eval=None, ... ) >>> datamodule.setup() @@ -121,10 +103,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST, @@ -135,10 +113,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, @@ -153,14 +127,12 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = BTechDataset( task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.root, category=self.category, ) self.test_data = BTechDataset( task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.root, category=self.category, diff --git a/src/anomalib/data/datamodules/image/folder.py b/src/anomalib/data/datamodules/image/folder.py index 7941ba2f7b..7fe51c32a0 100644 --- a/src/anomalib/data/datamodules/image/folder.py +++ b/src/anomalib/data/datamodules/image/folder.py @@ -9,8 +9,6 @@ from collections.abc import Sequence from pathlib import Path -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.image.folder import FolderDataset @@ -47,14 +45,6 @@ class Folder(AnomalibDataModule): Defaults to ``8``. task (TaskType, optional): Task type. Could be ``classification``, ``detection`` or ``segmentation``. Defaults to ``segmentation``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR``. test_split_ratio (float): Fraction of images from the train set that will be reserved for testing. @@ -102,8 +92,6 @@ class Folder(AnomalibDataModule): abnormal_dir="crack", task=TaskType.SEGMENTATION, mask_dir=dataset_root / "mask" / "crack", - image_size=256, - normalization=InputNormalizationMethod.NONE, ) folder_datamodule.setup() @@ -136,10 +124,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.FROM_TEST, @@ -164,10 +148,6 @@ def __init__( test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, val_split_ratio=val_split_ratio, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, seed=seed, ) @@ -186,7 +166,6 @@ def _setup(self, _stage: str | None = None) -> None: self.train_data = FolderDataset( name=self.name, task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.root, normal_dir=self.normal_dir, @@ -199,7 +178,6 @@ def _setup(self, _stage: str | None = None) -> None: self.test_data = FolderDataset( name=self.name, task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.root, normal_dir=self.normal_dir, diff --git a/src/anomalib/data/datamodules/image/kolektor.py b/src/anomalib/data/datamodules/image/kolektor.py index 2f8dc3b92b..c962e4fba7 100644 --- a/src/anomalib/data/datamodules/image/kolektor.py +++ b/src/anomalib/data/datamodules/image/kolektor.py @@ -20,18 +20,10 @@ import logging from pathlib import Path -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.image.kolektor import KolektorDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - TestSplitMode, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, TestSplitMode, ValSplitMode, download_and_extract logger = logging.getLogger(__name__) @@ -56,14 +48,6 @@ class Kolektor(AnomalibDataModule): Defaults to ``8``. task TaskType): Task type, 'classification', 'detection' or 'segmentation' Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR`` test_split_ratio (float): Fraction of images from the train set that will be reserved for testing. @@ -83,10 +67,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST, @@ -97,10 +77,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, @@ -114,13 +90,11 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = KolektorDataset( task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.root, ) self.test_data = KolektorDataset( task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.root, ) diff --git a/src/anomalib/data/datamodules/image/mvtec.py b/src/anomalib/data/datamodules/image/mvtec.py index 508a582380..a465ef52c1 100644 --- a/src/anomalib/data/datamodules/image/mvtec.py +++ b/src/anomalib/data/datamodules/image/mvtec.py @@ -28,18 +28,10 @@ import logging from pathlib import Path -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.image.mvtec import MVTecDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - TestSplitMode, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, TestSplitMode, ValSplitMode, download_and_extract logger = logging.getLogger(__name__) @@ -68,14 +60,6 @@ class MVTec(AnomalibDataModule): Defaults to ``8``. task TaskType): Task type, 'classification', 'detection' or 'segmentation' Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR``. test_split_ratio (float): Fraction of images from the train set that will be reserved for testing. @@ -103,10 +87,6 @@ class MVTec(AnomalibDataModule): >>> datamodule = MVTec(category="cable") - To change the image and batch size: - - >>> datamodule = MVTec(image_size=(512, 512), train_batch_size=16, eval_batch_size=8) - MVTec AD dataset does not provide a validation set. If you would like to use a separate validation set, you can use the ``val_split_mode`` and ``val_split_ratio`` arguments to create a validation set. @@ -129,10 +109,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST, @@ -142,10 +118,6 @@ def __init__( super().__init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, num_workers=num_workers, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, @@ -172,14 +144,12 @@ def _setup(self, _stage: str | None = None) -> None: """ self.train_data = MVTecDataset( task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.root, category=self.category, ) self.test_data = MVTecDataset( task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.root, category=self.category, diff --git a/src/anomalib/data/datamodules/image/visa.py b/src/anomalib/data/datamodules/image/visa.py index 30bf945c73..a445349702 100644 --- a/src/anomalib/data/datamodules/image/visa.py +++ b/src/anomalib/data/datamodules/image/visa.py @@ -28,18 +28,11 @@ from pathlib import Path import cv2 -from torchvision.transforms.v2 import Transform from anomalib import TaskType from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.image.visa import VisaDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - TestSplitMode, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, TestSplitMode, ValSplitMode, download_and_extract logger = logging.getLogger(__name__) @@ -66,14 +59,6 @@ class Visa(AnomalibDataModule): Defaults to ``8``. task (TaskType): Task type, 'classification', 'detection' or 'segmentation' Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. test_split_mode (TestSplitMode): Setting that determines how the testing subset is obtained. Defaults to ``TestSplitMode.FROM_DIR``. test_split_ratio (float): Fraction of images from the train set that will be reserved for testing. @@ -94,10 +79,6 @@ def __init__( eval_batch_size: int = 32, num_workers: int = 8, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, test_split_mode: TestSplitMode | str = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode | str = ValSplitMode.SAME_AS_TEST, @@ -108,10 +89,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, @@ -127,14 +104,12 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = VisaDataset( task=self.task, - transform=self.train_transform, split=Split.TRAIN, root=self.split_root, category=self.category, ) self.test_data = VisaDataset( task=self.task, - transform=self.eval_transform, split=Split.TEST, root=self.split_root, category=self.category, diff --git a/src/anomalib/data/datamodules/video/avenue.py b/src/anomalib/data/datamodules/video/avenue.py index 8914475081..86d068e761 100644 --- a/src/anomalib/data/datamodules/video/avenue.py +++ b/src/anomalib/data/datamodules/video/avenue.py @@ -21,18 +21,12 @@ import cv2 import scipy.io -from torchvision.transforms.v2 import Transform from anomalib import TaskType from anomalib.data.datamodules.base.video import AnomalibVideoDataModule from anomalib.data.datasets.base.video import VideoTargetFrame from anomalib.data.datasets.video.avenue import AvenueDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, ValSplitMode, download_and_extract logger = logging.getLogger(__name__) @@ -64,14 +58,6 @@ class Avenue(AnomalibVideoDataModule): Defaults to ``VideoTargetFrame.LAST``. task (TaskType): Task type, 'classification', 'detection' or 'segmentation' Defaults to ``TaskType.SEGMENTATION``. - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. train_batch_size (int, optional): Training batch size. Defaults to ``32``. eval_batch_size (int, optional): Test batch size. @@ -141,10 +127,6 @@ def __init__( frames_between_clips: int = 1, target_frame: VideoTargetFrame | str = VideoTargetFrame.LAST, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, train_batch_size: int = 32, eval_batch_size: int = 32, num_workers: int = 8, @@ -156,10 +138,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, val_split_mode=val_split_mode, val_split_ratio=val_split_ratio, seed=seed, @@ -175,7 +153,6 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = AvenueDataset( task=self.task, - transform=self.train_transform, clip_length_in_frames=self.clip_length_in_frames, frames_between_clips=self.frames_between_clips, target_frame=self.target_frame, @@ -186,7 +163,6 @@ def _setup(self, _stage: str | None = None) -> None: self.test_data = AvenueDataset( task=self.task, - transform=self.eval_transform, clip_length_in_frames=self.clip_length_in_frames, frames_between_clips=self.frames_between_clips, target_frame=self.target_frame, diff --git a/src/anomalib/data/datamodules/video/shanghaitech.py b/src/anomalib/data/datamodules/video/shanghaitech.py index b474f09547..2b5c6f428c 100644 --- a/src/anomalib/data/datamodules/video/shanghaitech.py +++ b/src/anomalib/data/datamodules/video/shanghaitech.py @@ -20,18 +20,11 @@ from pathlib import Path from shutil import move -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.video import AnomalibVideoDataModule from anomalib.data.datasets.base.video import VideoTargetFrame from anomalib.data.datasets.video.shanghaitech import ShanghaiTechDataset -from anomalib.data.utils import ( - DownloadInfo, - Split, - ValSplitMode, - download_and_extract, -) +from anomalib.data.utils import DownloadInfo, Split, ValSplitMode, download_and_extract from anomalib.data.utils.video import convert_video logger = logging.getLogger(__name__) @@ -53,14 +46,6 @@ class ShanghaiTech(AnomalibVideoDataModule): frames_between_clips (int, optional): Number of frames between each consecutive video clip. target_frame (VideoTargetFrame): Specifies the target frame in the video clip, used for ground truth retrieval task TaskType): Task type, 'classification', 'detection' or 'segmentation' - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. train_batch_size (int, optional): Training batch size. Defaults to 32. eval_batch_size (int, optional): Test batch size. Defaults to 32. num_workers (int, optional): Number of workers. Defaults to 8. @@ -77,10 +62,6 @@ def __init__( frames_between_clips: int = 1, target_frame: VideoTargetFrame = VideoTargetFrame.LAST, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, train_batch_size: int = 32, eval_batch_size: int = 32, num_workers: int = 8, @@ -92,10 +73,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, val_split_mode=val_split_mode, val_split_ratio=val_split_ratio, seed=seed, @@ -112,7 +89,6 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = ShanghaiTechDataset( task=self.task, - transform=self.train_transform, clip_length_in_frames=self.clip_length_in_frames, frames_between_clips=self.frames_between_clips, target_frame=self.target_frame, @@ -123,7 +99,6 @@ def _setup(self, _stage: str | None = None) -> None: self.test_data = ShanghaiTechDataset( task=self.task, - transform=self.eval_transform, clip_length_in_frames=self.clip_length_in_frames, frames_between_clips=self.frames_between_clips, target_frame=self.target_frame, diff --git a/src/anomalib/data/datamodules/video/ucsd_ped.py b/src/anomalib/data/datamodules/video/ucsd_ped.py index 2dd480ef37..4743d17044 100644 --- a/src/anomalib/data/datamodules/video/ucsd_ped.py +++ b/src/anomalib/data/datamodules/video/ucsd_ped.py @@ -7,8 +7,6 @@ from pathlib import Path from shutil import move -from torchvision.transforms.v2 import Transform - from anomalib import TaskType from anomalib.data.datamodules.base.video import AnomalibVideoDataModule from anomalib.data.datasets.base.video import VideoTargetFrame @@ -34,14 +32,6 @@ class UCSDped(AnomalibVideoDataModule): frames_between_clips (int, optional): Number of frames between each consecutive video clip. target_frame (VideoTargetFrame): Specifies the target frame in the video clip, used for ground truth retrieval task (TaskType): Task type, 'classification', 'detection' or 'segmentation' - image_size (tuple[int, int], optional): Size to which input images should be resized. - Defaults to ``None``. - transform (Transform, optional): Transforms that should be applied to the input images. - Defaults to ``None``. - train_transform (Transform, optional): Transforms that should be applied to the input images during training. - Defaults to ``None``. - eval_transform (Transform, optional): Transforms that should be applied to the input images during evaluation. - Defaults to ``None``. train_batch_size (int, optional): Training batch size. Defaults to 32. eval_batch_size (int, optional): Test batch size. Defaults to 32. num_workers (int, optional): Number of workers. Defaults to 8. @@ -58,10 +48,6 @@ def __init__( frames_between_clips: int = 10, target_frame: VideoTargetFrame = VideoTargetFrame.LAST, task: TaskType | str = TaskType.SEGMENTATION, - image_size: tuple[int, int] | None = None, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, train_batch_size: int = 8, eval_batch_size: int = 8, num_workers: int = 8, @@ -73,10 +59,6 @@ def __init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, val_split_mode=val_split_mode, val_split_ratio=val_split_ratio, seed=seed, @@ -93,7 +75,6 @@ def __init__( def _setup(self, _stage: str | None = None) -> None: self.train_data = UCSDpedDataset( task=self.task, - transform=self.train_transform, clip_length_in_frames=self.clip_length_in_frames, frames_between_clips=self.frames_between_clips, target_frame=self.target_frame, @@ -104,7 +85,6 @@ def _setup(self, _stage: str | None = None) -> None: self.test_data = UCSDpedDataset( task=self.task, - transform=self.eval_transform, clip_length_in_frames=self.clip_length_in_frames, frames_between_clips=self.frames_between_clips, target_frame=self.target_frame, diff --git a/src/anomalib/deploy/utils.py b/src/anomalib/deploy/utils.py deleted file mode 100644 index e2f23bf841..0000000000 --- a/src/anomalib/deploy/utils.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Utility functions for Anomalib deployment module.""" - -# Copyright (C) 2024 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -from torchvision.transforms.v2 import CenterCrop, Compose, Resize, Transform - -from anomalib.data.transforms import ExportableCenterCrop - - -def make_transform_exportable(transform: Transform) -> Transform: - """Get exportable transform. - - Some transforms are not supported by ONNX/OpenVINO, so we need to replace them with exportable versions. - """ - transform = disable_antialiasing(transform) - return convert_centercrop(transform) - - -def disable_antialiasing(transform: Transform) -> Transform: - """Disable antialiasing in Resize transforms. - - Resizing with antialiasing is not supported by ONNX, so we need to disable it. - """ - if isinstance(transform, Resize): - transform.antialias = False - if isinstance(transform, Compose): - for tr in transform.transforms: - disable_antialiasing(tr) - return transform - - -def convert_centercrop(transform: Transform) -> Transform: - """Convert CenterCrop to ExportableCenterCrop. - - Torchvision's CenterCrop is not supported by ONNX, so we need to replace it with our own ExportableCenterCrop. - """ - if isinstance(transform, CenterCrop): - transform = ExportableCenterCrop(size=transform.size) - if isinstance(transform, Compose): - for index in range(len(transform.transforms)): - tr = transform.transforms[index] - transform.transforms[index] = convert_centercrop(tr) - return transform diff --git a/src/anomalib/engine/engine.py b/src/anomalib/engine/engine.py index 13eef8a63c..36bfcc3bf4 100644 --- a/src/anomalib/engine/engine.py +++ b/src/anomalib/engine/engine.py @@ -8,7 +8,6 @@ from pathlib import Path from typing import Any -import torch from lightning.pytorch.callbacks import Callback from lightning.pytorch.loggers import Logger from lightning.pytorch.trainer import Trainer @@ -292,60 +291,6 @@ def _setup_dataset_task( ) data.task = self.task - @staticmethod - def _setup_transform( - model: AnomalyModule, - datamodule: AnomalibDataModule | None = None, - dataloaders: EVAL_DATALOADERS | TRAIN_DATALOADERS | None = None, - ckpt_path: Path | str | None = None, - ) -> None: - """Implements the logic for setting the transform at the start of each run. - - Any transform passed explicitly to the datamodule takes precedence. Otherwise, if a checkpoint path is provided, - we can load the transform from the checkpoint. If no transform is provided, we use the default transform from - the model. - - Args: - model (AnomalyModule): The model to assign the transform to. - datamodule (AnomalibDataModule | None): The datamodule to assign the transform from. - defaults to ``None``. - dataloaders (EVAL_DATALOADERS | TRAIN_DATALOADERS | None): Dataloaders to assign the transform to. - defaults to ``None``. - ckpt_path (str): The path to the checkpoint. - defaults to ``None``. - - Returns: - Transform: The transform loaded from the checkpoint. - """ - if isinstance(dataloaders, DataLoader): - dataloaders = [dataloaders] - - # get transform - if datamodule and datamodule.transform: - # a transform passed explicitly to the datamodule takes precedence - transform = datamodule.transform - elif dataloaders and any(getattr(dl.dataset, "transform", None) for dl in dataloaders): - # if dataloaders are provided, we use the transform from the first dataloader that has a transform - transform = next(dl.dataset.transform for dl in dataloaders if getattr(dl.dataset, "transform", None)) - elif ckpt_path is not None: - # if a checkpoint path is provided, we can load the transform from the checkpoint - checkpoint = torch.load(ckpt_path, map_location=model.device) - transform = checkpoint["transform"] - elif model.transform is None: - # if no transform is provided, we use the default transform from the model - image_size = datamodule.image_size if datamodule else None - transform = model.configure_transforms(image_size) - else: - transform = model.transform - - # update transform in model - model.set_transform(transform) - # The dataloaders don't have access to the trainer and/or model, so we need to set the transforms manually - if dataloaders: - for dataloader in dataloaders: - if not getattr(dataloader.dataset, "transform", None): - dataloader.dataset.transform = transform - def _setup_anomalib_callbacks(self, model: AnomalyModule) -> None: """Set up callbacks for the trainer.""" _callbacks: list[Callback] = [] @@ -458,7 +403,6 @@ def fit( ) self._setup_trainer(model) self._setup_dataset_task(train_dataloaders, val_dataloaders, datamodule) - self._setup_transform(model, datamodule=datamodule, ckpt_path=ckpt_path) if model.learning_type in {LearningType.ZERO_SHOT, LearningType.FEW_SHOT}: # if the model is zero-shot or few-shot, we only need to run validate for normalization and thresholding self.trainer.validate(model, val_dataloaders, datamodule=datamodule, ckpt_path=ckpt_path) @@ -512,7 +456,6 @@ def validate( if model: self._setup_trainer(model) self._setup_dataset_task(dataloaders) - self._setup_transform(model or self.model, datamodule=datamodule, ckpt_path=ckpt_path) return self.trainer.validate(model, dataloaders, ckpt_path, verbose, datamodule) def test( @@ -606,7 +549,6 @@ def test( raise RuntimeError(msg) self._setup_dataset_task(dataloaders) - self._setup_transform(model or self.model, datamodule=datamodule, ckpt_path=ckpt_path) if self._should_run_validation(model or self.model, ckpt_path): logger.info("Running validation before testing to collect normalization metrics and/or thresholds.") self.trainer.validate(model, dataloaders, None, verbose=False, datamodule=datamodule) @@ -711,7 +653,6 @@ def predict( dataloaders = dataloaders or None self._setup_dataset_task(dataloaders, datamodule) - self._setup_transform(model or self.model, datamodule=datamodule, dataloaders=dataloaders, ckpt_path=ckpt_path) if self._should_run_validation(model or self.model, ckpt_path): logger.info("Running validation before predicting to collect normalization metrics and/or thresholds.") @@ -781,7 +722,6 @@ def train( test_dataloaders, datamodule, ) - self._setup_transform(model, datamodule=datamodule, ckpt_path=ckpt_path) if model.learning_type in {LearningType.ZERO_SHOT, LearningType.FEW_SHOT}: # if the model is zero-shot or few-shot, we only need to run validate for normalization and thresholding self.trainer.validate(model, val_dataloaders, None, verbose=False, datamodule=datamodule) @@ -828,8 +768,7 @@ def export( Path: Path to the exported model. Raises: - ValueError: If Dataset, Datamodule, and transform are not provided. - TypeError: If path to the transform file is not a string or Path. + ValueError: If Dataset, Datamodule are not provided. CLI Usage: 1. To export as a torch ``.pt`` file you can run the following command. diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py index b22ee6981b..ff12db0cec 100644 --- a/src/anomalib/models/components/base/anomaly_module.py +++ b/src/anomalib/models/components/base/anomaly_module.py @@ -5,8 +5,9 @@ import logging from abc import ABC, abstractmethod +from collections.abc import Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any import lightning.pytorch as pl import torch @@ -14,7 +15,7 @@ from lightning.pytorch.trainer.states import TrainerFn from lightning.pytorch.utilities.types import STEP_OUTPUT from torch import nn -from torchvision.transforms.v2 import Compose, Normalize, Resize, Transform +from torchvision.transforms.v2 import Compose, Normalize, Resize from anomalib import LearningType from anomalib.data import Batch, InferenceBatch @@ -22,13 +23,10 @@ from anomalib.metrics.evaluator import Evaluator from anomalib.metrics.threshold import Threshold from anomalib.post_processing import OneClassPostProcessor, PostProcessor +from anomalib.pre_processing import PreProcessor from .export_mixin import ExportMixin -if TYPE_CHECKING: - from lightning.pytorch.callbacks import Callback - - logger = logging.getLogger(__name__) @@ -40,6 +38,7 @@ class AnomalyModule(ExportMixin, pl.LightningModule, ABC): def __init__( self, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: @@ -51,14 +50,11 @@ def __init__( self.loss: nn.Module self.callbacks: list[Callback] - # set the post-processor + self.pre_processor = self._resolve_pre_processor(pre_processor) self.post_processor = post_processor or self.default_post_processor() - self.evaluator = self._resolve_evaluator(evaluator) - self._transform: Transform | None = None self._input_size: tuple[int, int] | None = None - self._is_setup = False # flag to track if setup has been called from the trainer @property @@ -82,6 +78,29 @@ def _setup(self) -> None: initialization. """ + def _resolve_pre_processor(self, pre_processor: PreProcessor | bool) -> PreProcessor | None: + """Resolve and validate which pre-processor to use.. + + Args: + pre_processor: Pre-processor configuration + - True -> use default pre-processor + - False -> no pre-processor + - PreProcessor -> use the provided pre-processor + + Returns: + Configured pre-processor + """ + if isinstance(pre_processor, PreProcessor): + return pre_processor + if isinstance(pre_processor, bool): + return self.configure_pre_processor() if pre_processor else None + msg = f"Invalid pre-processor type: {type(pre_processor)}" + raise TypeError(msg) + + def configure_callbacks(self) -> Sequence[Callback] | Callback: + """Configure default callbacks for AnomalyModule.""" + return [self.pre_processor] if self.pre_processor else [] + def forward(self, batch: torch.Tensor, *args, **kwargs) -> InferenceBatch: """Perform the forward-pass by passing input tensor to the module. @@ -94,7 +113,7 @@ def forward(self, batch: torch.Tensor, *args, **kwargs) -> InferenceBatch: Tensor: Output tensor from the model. """ del args, kwargs # These variables are not used. - batch = self.exportable_transform(batch) + batch = self.pre_processor(batch) if self.pre_processor else batch batch = self.model(batch) return self.post_processor(batch) if self.post_processor else batch @@ -150,36 +169,48 @@ def learning_type(self) -> LearningType: """Learning type of the model.""" raise NotImplementedError - @property - def transform(self) -> Transform: - """Retrieve the model-specific transform. + @classmethod + def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> PreProcessor: + """Configure the pre-processor. - If a transform has been set using `set_transform`, it will be returned. Otherwise, we will use the - model-specific default transform, conditioned on the input size. - """ - return self._transform + The default pre-processor resizes images to 256x256 and normalizes using ImageNet statistics. + Individual models can override this method to provide custom transforms and pre-processing pipelines. - def set_transform(self, transform: Transform) -> None: - """Update the transform linked to the model instance.""" - self._transform = transform + Args: + image_size (tuple[int, int] | None, optional): Target size for resizing images. + If None, defaults to (256, 256). Defaults to None. + **kwargs (Any): Additional keyword arguments (unused). + + Returns: + PreProcessor: Configured pre-processor instance. + + Examples: + Get default pre-processor with custom image size: - def configure_transforms(self, image_size: tuple[int, int] | None = None) -> Transform: # noqa: PLR6301 - """Default transforms. + >>> preprocessor = AnomalyModule.configure_pre_processor(image_size=(512, 512)) - The default transform is resize to 256x256 and normalize to ImageNet stats. Individual models can override - this method to provide custom transforms. + Create model with custom pre-processor: + + >>> from torchvision.transforms.v2 import RandomHorizontalFlip + >>> custom_transform = Compose([ + ... Resize((256, 256), antialias=True), + ... CenterCrop((224, 224)), + ... RandomHorizontalFlip(p=0.5), + ... Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ... ]) + >>> preprocessor.train_transform = custom_transform + >>> model = PatchCore(pre_processor=preprocessor) + + Disable pre-processing: + + >>> model = PatchCore(pre_processor=False) """ - logger.warning( - "No implementation of `configure_transforms` was provided in the Lightning model. Using default " - "transforms from the base class. This may not be suitable for your use case. Please override " - "`configure_transforms` in your model.", - ) image_size = image_size or (256, 256) - return Compose( - [ + return PreProcessor( + transform=Compose([ Resize(image_size, antialias=True), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ], + ]), ) def default_post_processor(self) -> PostProcessor: @@ -226,30 +257,12 @@ def input_size(self) -> tuple[int, int] | None: The effective input size is the size of the input tensor after the transform has been applied. If the transform is not set, or if the transform does not change the shape of the input tensor, this method will return None. """ - transform = self.transform or self.configure_transforms() + transform = self.pre_processor.predict_transform if self.pre_processor else None if transform is None: return None dummy_input = torch.zeros(1, 3, 1, 1) output_shape = transform(dummy_input).shape[-2:] - if output_shape == (1, 1): - return None - return output_shape[-2:] - - def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None: - """Called when saving the model to a checkpoint. - - Saves the transform to the checkpoint. - """ - checkpoint["transform"] = self.transform - - def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None: - """Called when loading the model from a checkpoint. - - Loads the transform from the checkpoint and calls setup to ensure that the torch model is built before loading - the state dict. - """ - self._transform = checkpoint["transform"] - self.setup("load_checkpoint") + return None if output_shape == (1, 1) else output_shape[-2:] @classmethod def from_config( diff --git a/src/anomalib/models/components/base/export_mixin.py b/src/anomalib/models/components/base/export_mixin.py index b696ad2567..a0f84d1510 100644 --- a/src/anomalib/models/components/base/export_mixin.py +++ b/src/anomalib/models/components/base/export_mixin.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import logging -from collections.abc import Callable, Iterable +from collections.abc import Iterable from pathlib import Path from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any @@ -14,12 +14,10 @@ from lightning_utilities.core.imports import package_available from torch import nn from torchmetrics import Metric -from torchvision.transforms.v2 import Transform from anomalib import TaskType from anomalib.data import AnomalibDataModule from anomalib.deploy.export import CompressionType, ExportType -from anomalib.deploy.utils import make_transform_exportable if TYPE_CHECKING: from importlib.util import find_spec @@ -34,8 +32,6 @@ class ExportMixin: """This mixin allows exporting models to torch and ONNX/OpenVINO.""" model: nn.Module - transform: Transform - configure_transforms: Callable device: torch.device def to_torch( @@ -136,7 +132,7 @@ def to_onnx( dynamic_axes = ( {"input": {0: "batch_size"}, "output": {0: "batch_size"}} if input_size - else {"input": {0: "batch_size", 2: "height", 3: "weight"}, "output": {0: "batch_size"}} + else {"input": {0: "batch_size", 2: "height", 3: "width"}, "output": {0: "batch_size"}} ) onnx_path = export_root / "model.onnx" # apply pass through the model to get the output names @@ -400,11 +396,6 @@ def val_fn(nncf_model: "CompiledModel", validation_data: Iterable) -> float: return nncf.quantize_with_accuracy_control(model, calibration_dataset, validation_dataset, val_fn) - @property - def exportable_transform(self) -> Transform: - """Return the exportable transform.""" - return make_transform_exportable(self.transform) - def _create_export_root(export_root: str | Path, export_type: ExportType) -> Path: """Create export directory. diff --git a/src/anomalib/models/image/cfa/lightning_model.py b/src/anomalib/models/image/cfa/lightning_model.py index 363bd2eae7..154ea4e3e8 100644 --- a/src/anomalib/models/image/cfa/lightning_model.py +++ b/src/anomalib/models/image/cfa/lightning_model.py @@ -19,6 +19,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import CfaLoss from .torch_model import CfaModel @@ -44,6 +45,9 @@ class Cfa(AnomalyModule): Defaults to ``3``. radius (float): Radius of the hypersphere to search the soft boundary. Defaults to ``1e-5``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -54,10 +58,11 @@ def __init__( num_nearest_neighbors: int = 3, num_hard_negative_features: int = 3, radius: float = 1e-5, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: CfaModel = CfaModel( backbone=backbone, gamma_c=gamma_c, diff --git a/src/anomalib/models/image/cflow/lightning_model.py b/src/anomalib/models/image/cflow/lightning_model.py index 3b4cb731e2..b6118d8e4e 100644 --- a/src/anomalib/models/image/cflow/lightning_model.py +++ b/src/anomalib/models/image/cflow/lightning_model.py @@ -26,6 +26,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import CflowModel from .utils import get_logp, positional_encoding_2d @@ -69,10 +70,11 @@ def __init__( clamp_alpha: float = 1.9, permute_soft: bool = False, lr: float = 0.0001, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: CflowModel = CflowModel( backbone=backbone, diff --git a/src/anomalib/models/image/csflow/lightning_model.py b/src/anomalib/models/image/csflow/lightning_model.py index e3762da180..0ae381c65b 100644 --- a/src/anomalib/models/image/csflow/lightning_model.py +++ b/src/anomalib/models/image/csflow/lightning_model.py @@ -17,6 +17,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import CsFlowLoss from .torch_model import CsFlowModel @@ -46,25 +47,20 @@ def __init__( n_coupling_blocks: int = 4, clamp: int = 3, num_channels: int = 3, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) + if self.input_size is None: + msg = "CsFlow needs input size to build torch model." + raise ValueError(msg) self.cross_conv_hidden_channels = cross_conv_hidden_channels self.n_coupling_blocks = n_coupling_blocks self.clamp = clamp self.num_channels = num_channels - self.loss = CsFlowLoss() - - self.model: CsFlowModel - - def _setup(self) -> None: - if self.input_size is None: - msg = "CsFlow needs input size to build torch model." - raise ValueError(msg) - self.model = CsFlowModel( input_size=self.input_size, cross_conv_hidden_channels=self.cross_conv_hidden_channels, @@ -73,6 +69,7 @@ def _setup(self) -> None: num_channels=self.num_channels, ) self.model.feature_extractor.eval() + self.loss = CsFlowLoss() def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: """Perform the training step of CS-Flow. diff --git a/src/anomalib/models/image/dfkde/lightning_model.py b/src/anomalib/models/image/dfkde/lightning_model.py index 8b67e56907..9bd8388d49 100644 --- a/src/anomalib/models/image/dfkde/lightning_model.py +++ b/src/anomalib/models/image/dfkde/lightning_model.py @@ -16,6 +16,7 @@ from anomalib.models.components import AnomalyModule, MemoryBankMixin from anomalib.models.components.classification import FeatureScalingMethod from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import DfkdeModel @@ -48,10 +49,11 @@ def __init__( n_pca_components: int = 16, feature_scaling_method: FeatureScalingMethod = FeatureScalingMethod.SCALE, max_training_points: int = 40000, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model = DfkdeModel( layers=layers, diff --git a/src/anomalib/models/image/dfm/lightning_model.py b/src/anomalib/models/image/dfm/lightning_model.py index 9b6f52979c..b0449d1e69 100644 --- a/src/anomalib/models/image/dfm/lightning_model.py +++ b/src/anomalib/models/image/dfm/lightning_model.py @@ -17,6 +17,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule, MemoryBankMixin from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import DFMModel @@ -39,6 +40,9 @@ class Dfm(MemoryBankMixin, AnomalyModule): Defaults to ``0.97``. score_type (str, optional): Scoring type. Options are `fre` and `nll`. Defaults to ``fre``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -49,10 +53,11 @@ def __init__( pooling_kernel_size: int = 4, pca_level: float = 0.97, score_type: str = "fre", + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: DFMModel = DFMModel( backbone=backbone, diff --git a/src/anomalib/models/image/draem/lightning_model.py b/src/anomalib/models/image/draem/lightning_model.py index ccfb52cbbd..a072bcae0f 100644 --- a/src/anomalib/models/image/draem/lightning_model.py +++ b/src/anomalib/models/image/draem/lightning_model.py @@ -20,6 +20,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import DraemLoss from .torch_model import DraemModel @@ -38,6 +39,9 @@ class Draem(AnomalyModule): anomaly_source_path (str | None): Path to folder that contains the anomaly source images. Random noise will be used if left empty. Defaults to ``None``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -46,10 +50,11 @@ def __init__( sspcab_lambda: float = 0.1, anomaly_source_path: str | None = None, beta: float | tuple[float, float] = (0.1, 1.0), + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.augmenter = Augmenter(anomaly_source_path, beta=beta) self.model = DraemModel(sspcab=enable_sspcab) diff --git a/src/anomalib/models/image/dsr/lightning_model.py b/src/anomalib/models/image/dsr/lightning_model.py index 8ae8633c9c..a4ed2df231 100644 --- a/src/anomalib/models/image/dsr/lightning_model.py +++ b/src/anomalib/models/image/dsr/lightning_model.py @@ -24,6 +24,7 @@ from anomalib.models.image.dsr.loss import DsrSecondStageLoss, DsrThirdStageLoss from anomalib.models.image.dsr.torch_model import DsrModel from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor __all__ = ["Dsr"] @@ -42,16 +43,20 @@ class Dsr(AnomalyModule): Args: latent_anomaly_strength (float): Strength of the generated anomalies in the latent space. Defaults to 0.2 upsampling_train_ratio (float): Ratio of training steps for the upsampling module. Defaults to 0.7 + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( self, latent_anomaly_strength: float = 0.2, upsampling_train_ratio: float = 0.7, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.automatic_optimization = False self.upsampling_train_ratio = upsampling_train_ratio diff --git a/src/anomalib/models/image/efficient_ad/lightning_model.py b/src/anomalib/models/image/efficient_ad/lightning_model.py index 152f50c36a..88b29f7215 100644 --- a/src/anomalib/models/image/efficient_ad/lightning_model.py +++ b/src/anomalib/models/image/efficient_ad/lightning_model.py @@ -15,7 +15,7 @@ from lightning.pytorch.utilities.types import STEP_OUTPUT from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder -from torchvision.transforms.v2 import CenterCrop, Compose, Normalize, RandomGrayscale, Resize, ToTensor, Transform +from torchvision.transforms.v2 import CenterCrop, Compose, Normalize, RandomGrayscale, Resize, ToTensor from anomalib import LearningType from anomalib.data import Batch @@ -23,6 +23,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import EfficientAdModel, EfficientAdModelSize, reduce_tensor_elems @@ -60,6 +61,9 @@ class EfficientAd(AnomalyModule): pad_maps (bool): relevant if padding is set to False. In this case, pad_maps = True pads the output anomaly maps so that their size matches the size in the padding = True case. Defaults to ``True``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -71,10 +75,11 @@ def __init__( weight_decay: float = 0.00001, padding: bool = False, pad_maps: bool = True, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.imagenet_dir = Path(imagenet_dir) if not isinstance(model_size, EfficientAdModelSize): @@ -207,6 +212,13 @@ def _get_quantiles_of_maps(self, maps: list[torch.Tensor]) -> tuple[torch.Tensor qb = torch.quantile(maps_flat, q=0.995).to(self.device) return qa, qb + @classmethod + def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> PreProcessor: + """Default transform for EfficientAd. Imagenet normalization applied in forward.""" + image_size = image_size or (256, 256) + transform = Compose([Resize(image_size, antialias=True)]) + return PreProcessor(transform=transform) + def configure_optimizers(self) -> torch.optim.Optimizer: """Configure optimizers.""" optimizer = torch.optim.Adam( @@ -247,9 +259,12 @@ def on_train_start(self) -> None: if self.trainer.datamodule.train_batch_size != 1: msg = "train_batch_size for EfficientAd should be 1." raise ValueError(msg) - if self._transform and any(isinstance(transform, Normalize) for transform in self._transform.transforms): - msg = "Transforms for EfficientAd should not contain Normalize." - raise ValueError(msg) + + if self.pre_processor and self.pre_processor.train_transform: + transforms = self.pre_processor.train_transform.transforms + if transforms and any(isinstance(transform, Normalize) for transform in transforms): + msg = "Transforms for EfficientAd should not contain Normalize." + raise ValueError(msg) sample = next(iter(self.trainer.train_dataloader)) image_size = sample.image.shape[-2:] @@ -322,13 +337,3 @@ def learning_type(self) -> LearningType: LearningType: Learning type of the model. """ return LearningType.ONE_CLASS - - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform: - """Default transform for EfficientAd. Imagenet normalization applied in forward.""" - image_size = image_size or (256, 256) - return Compose( - [ - Resize(image_size, antialias=True), - ], - ) diff --git a/src/anomalib/models/image/fastflow/lightning_model.py b/src/anomalib/models/image/fastflow/lightning_model.py index 75aff99584..935df8468d 100644 --- a/src/anomalib/models/image/fastflow/lightning_model.py +++ b/src/anomalib/models/image/fastflow/lightning_model.py @@ -17,6 +17,7 @@ from anomalib.metrics import AUROC, Evaluator, F1Score from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import FastflowLoss from .torch_model import FastflowModel @@ -35,7 +36,10 @@ class Fastflow(AnomalyModule): conv3x3_only (bool, optinoal): Use only conv3x3 in fast_flow model. Defaults to ``False``. hidden_ratio (float, optional): Ratio to calculate hidden var channels. - Defaults to ``1.0`. + Defaults to ``1.0``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -45,10 +49,14 @@ def __init__( flow_steps: int = 8, conv3x3_only: bool = False, hidden_ratio: float = 1.0, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) + if self.input_size is None: + msg = "Fastflow needs input size to build torch model." + raise ValueError(msg) self.backbone = backbone self.pre_trained = pre_trained @@ -56,14 +64,6 @@ def __init__( self.conv3x3_only = conv3x3_only self.hidden_ratio = hidden_ratio - self.model: FastflowModel - self.loss = FastflowLoss() - - def _setup(self) -> None: - if self.input_size is None: - msg = "Fastflow needs input size to build torch model." - raise ValueError(msg) - self.model = FastflowModel( input_size=self.input_size, backbone=self.backbone, @@ -72,6 +72,7 @@ def _setup(self) -> None: conv3x3_only=self.conv3x3_only, hidden_ratio=self.hidden_ratio, ) + self.loss = FastflowLoss() def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: """Perform the training step input and return the loss. diff --git a/src/anomalib/models/image/fre/lightning_model.py b/src/anomalib/models/image/fre/lightning_model.py index b4628e6446..f3de232667 100755 --- a/src/anomalib/models/image/fre/lightning_model.py +++ b/src/anomalib/models/image/fre/lightning_model.py @@ -18,6 +18,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import FREModel @@ -41,6 +42,9 @@ class Fre(AnomalyModule): latent_dim (int, optional): Reduced size of feature after applying dimensionality reduction via shallow linear autoencoder. Defaults to ``220``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -51,10 +55,11 @@ def __init__( pooling_kernel_size: int = 2, input_dim: int = 65536, latent_dim: int = 220, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: FREModel = FREModel( backbone=backbone, diff --git a/src/anomalib/models/image/ganomaly/lightning_model.py b/src/anomalib/models/image/ganomaly/lightning_model.py index 362a713050..de3a479aa8 100644 --- a/src/anomalib/models/image/ganomaly/lightning_model.py +++ b/src/anomalib/models/image/ganomaly/lightning_model.py @@ -18,6 +18,7 @@ from anomalib.metrics import AUROC, Evaluator, F1Score from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import DiscriminatorLoss, GeneratorLoss from .torch_model import GanomalyModel @@ -51,6 +52,9 @@ class Ganomaly(AnomalyModule): Defaults to ``0.5``. beta2 (float, optional): Adam beta2. Defaults to ``0.999``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -66,10 +70,14 @@ def __init__( lr: float = 0.0002, beta1: float = 0.5, beta2: float = 0.999, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) + if self.input_size is None: + msg = "GANomaly needs input size to build torch model." + raise ValueError(msg) self.n_features = n_features self.latent_vec_size = latent_vec_size @@ -82,6 +90,15 @@ def __init__( self.min_scores: torch.Tensor = torch.tensor(float("inf"), dtype=torch.float32) # pylint: disable=not-callable self.max_scores: torch.Tensor = torch.tensor(float("-inf"), dtype=torch.float32) # pylint: disable=not-callable + self.model = GanomalyModel( + input_size=self.input_size, + num_input_channels=3, + n_features=self.n_features, + latent_vec_size=self.latent_vec_size, + extra_layers=self.extra_layers, + add_final_conv_layer=self.add_final_conv_layer, + ) + self.generator_loss = GeneratorLoss(wadv, wcon, wenc) self.discriminator_loss = DiscriminatorLoss() self.automatic_optimization = False @@ -94,20 +111,6 @@ def __init__( self.model: GanomalyModel - def _setup(self) -> None: - if self.input_size is None: - msg = "GANomaly needs input size to build torch model." - raise ValueError(msg) - - self.model = GanomalyModel( - input_size=self.input_size, - num_input_channels=3, - n_features=self.n_features, - latent_vec_size=self.latent_vec_size, - extra_layers=self.extra_layers, - add_final_conv_layer=self.add_final_conv_layer, - ) - def _reset_min_max(self) -> None: """Reset min_max scores.""" self.min_scores = torch.tensor(float("inf"), dtype=torch.float32) # pylint: disable=not-callable diff --git a/src/anomalib/models/image/padim/lightning_model.py b/src/anomalib/models/image/padim/lightning_model.py index 2a80171931..aed2163def 100644 --- a/src/anomalib/models/image/padim/lightning_model.py +++ b/src/anomalib/models/image/padim/lightning_model.py @@ -10,14 +10,13 @@ import torch from lightning.pytorch.utilities.types import STEP_OUTPUT -from torchvision.transforms.v2 import Compose, Normalize, Resize, Transform from anomalib import LearningType from anomalib.data import Batch from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule, MemoryBankMixin -from anomalib.post_processing import PostProcessor -from anomalib.post_processing.one_class import OneClassPostProcessor +from anomalib.post_processing import OneClassPostProcessor, PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import PadimModel @@ -39,6 +38,9 @@ class Padim(MemoryBankMixin, AnomalyModule): n_features (int, optional): Number of features to retain in the dimension reduction step. Default values from the paper are available for: resnet18 (100), wide_resnet50_2 (550). Defaults to ``None``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -47,10 +49,11 @@ def __init__( layers: list[str] = ["layer1", "layer2", "layer3"], # noqa: B006 pre_trained: bool = True, n_features: int | None = None, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: PadimModel = PadimModel( backbone=backbone, @@ -128,17 +131,6 @@ def learning_type(self) -> LearningType: """ return LearningType.ONE_CLASS - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform: - """Default transform for Padim.""" - image_size = image_size or (256, 256) - return Compose( - [ - Resize(image_size, antialias=True), - Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ], - ) - @staticmethod def default_post_processor() -> OneClassPostProcessor: """Return the default post-processor for PADIM.""" diff --git a/src/anomalib/models/image/patchcore/lightning_model.py b/src/anomalib/models/image/patchcore/lightning_model.py index d633141af3..f855a61d8e 100644 --- a/src/anomalib/models/image/patchcore/lightning_model.py +++ b/src/anomalib/models/image/patchcore/lightning_model.py @@ -12,14 +12,14 @@ import torch from lightning.pytorch.utilities.types import STEP_OUTPUT -from torchvision.transforms.v2 import CenterCrop, Compose, Normalize, Resize, Transform +from torchvision.transforms.v2 import CenterCrop, Compose, Normalize, Resize from anomalib import LearningType from anomalib.data import Batch from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule, MemoryBankMixin -from anomalib.post_processing import PostProcessor -from anomalib.post_processing.one_class import OneClassPostProcessor +from anomalib.post_processing import OneClassPostProcessor, PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import PatchcoreModel @@ -40,6 +40,9 @@ class Patchcore(MemoryBankMixin, AnomalyModule): Defaults to ``0.1``. num_neighbors (int, optional): Number of nearest neighbors. Defaults to ``9``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -49,10 +52,11 @@ def __init__( pre_trained: bool = True, coreset_sampling_ratio: float = 0.1, num_neighbors: int = 9, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: PatchcoreModel = PatchcoreModel( backbone=backbone, @@ -63,6 +67,26 @@ def __init__( self.coreset_sampling_ratio = coreset_sampling_ratio self.embeddings: list[torch.Tensor] = [] + @classmethod + def configure_pre_processor( + cls, + image_size: tuple[int, int] | None = None, + center_crop_size: tuple[int, int] | None = None, + ) -> PreProcessor: + """Default transform for Padim.""" + image_size = image_size or (256, 256) + if center_crop_size is None: + # scale center crop size proportional to image size + height, width = image_size + center_crop_size = (int(height * (224 / 256)), int(width * (224 / 256))) + + transform = Compose([ + Resize(image_size, antialias=True), + CenterCrop(center_crop_size), + Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + return PreProcessor(transform=transform) + @staticmethod def configure_optimizers() -> None: """Configure optimizers. @@ -129,21 +153,6 @@ def learning_type(self) -> LearningType: """ return LearningType.ONE_CLASS - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform: - """Default transform for Padim.""" - image_size = image_size or (256, 256) - # scale center crop size proportional to image size - height, width = image_size - center_crop_size = (int(height * (224 / 256)), int(width * (224 / 256))) - return Compose( - [ - Resize(image_size, antialias=True), - CenterCrop(center_crop_size), - Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ], - ) - @staticmethod def default_post_processor() -> OneClassPostProcessor: """Return the default post-processor for the model. diff --git a/src/anomalib/models/image/reverse_distillation/lightning_model.py b/src/anomalib/models/image/reverse_distillation/lightning_model.py index 916541ebda..a0cd8521f9 100644 --- a/src/anomalib/models/image/reverse_distillation/lightning_model.py +++ b/src/anomalib/models/image/reverse_distillation/lightning_model.py @@ -17,6 +17,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .anomaly_map import AnomalyMapGenerationMode from .loss import ReverseDistillationLoss @@ -35,6 +36,9 @@ class ReverseDistillation(AnomalyModule): Defaults to ``AnomalyMapGenerationMode.ADD``. pre_trained (bool, optional): Boolean to check whether to use a pre_trained backbone. Defaults to ``True``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -43,24 +47,20 @@ def __init__( layers: Sequence[str] = ("layer1", "layer2", "layer3"), anomaly_map_mode: AnomalyMapGenerationMode = AnomalyMapGenerationMode.ADD, pre_trained: bool = True, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) + if self.input_size is None: + msg = "Input size is required for Reverse Distillation model." + raise ValueError(msg) self.backbone = backbone self.pre_trained = pre_trained self.layers = layers self.anomaly_map_mode = anomaly_map_mode - self.model: ReverseDistillationModel - self.loss = ReverseDistillationLoss() - - def _setup(self) -> None: - if self.input_size is None: - msg = "Input size is required for Reverse Distillation model." - raise ValueError(msg) - self.model = ReverseDistillationModel( backbone=self.backbone, pre_trained=self.pre_trained, @@ -68,6 +68,7 @@ def _setup(self) -> None: input_size=self.input_size, anomaly_map_mode=self.anomaly_map_mode, ) + self.loss = ReverseDistillationLoss() def configure_optimizers(self) -> optim.Adam: """Configure optimizers for decoder and bottleneck. diff --git a/src/anomalib/models/image/rkde/lightning_model.py b/src/anomalib/models/image/rkde/lightning_model.py index 2d48da35c4..8d618127c0 100644 --- a/src/anomalib/models/image/rkde/lightning_model.py +++ b/src/anomalib/models/image/rkde/lightning_model.py @@ -18,6 +18,7 @@ from anomalib.models.components import AnomalyModule, MemoryBankMixin from anomalib.models.components.classification import FeatureScalingMethod from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .region_extractor import RoiStage from .torch_model import RkdeModel @@ -47,6 +48,9 @@ class Rkde(MemoryBankMixin, AnomalyModule): Defaults to ``FeatureScalingMethod.SCALE``. max_training_points (int, optional): Maximum number of training points to fit the KDE model. Defaults to ``40000``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -59,10 +63,11 @@ def __init__( n_pca_components: int = 16, feature_scaling_method: FeatureScalingMethod = FeatureScalingMethod.SCALE, max_training_points: int = 40000, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) self.model: RkdeModel = RkdeModel( roi_stage=roi_stage, diff --git a/src/anomalib/models/image/stfpm/lightning_model.py b/src/anomalib/models/image/stfpm/lightning_model.py index eca57f6850..32cace71f1 100644 --- a/src/anomalib/models/image/stfpm/lightning_model.py +++ b/src/anomalib/models/image/stfpm/lightning_model.py @@ -18,6 +18,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import STFPMLoss from .torch_model import STFPMModel @@ -33,21 +34,22 @@ class Stfpm(AnomalyModule): Defaults to ``resnet18``. layers (list[str]): Layers to extract features from the backbone CNN Defaults to ``["layer1", "layer2", "layer3"]``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( self, backbone: str = "resnet18", layers: Sequence[str] = ("layer1", "layer2", "layer3"), + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) - self.model = STFPMModel( - backbone=backbone, - layers=layers, - ) + self.model = STFPMModel(backbone=backbone, layers=layers) self.loss = STFPMLoss() def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: diff --git a/src/anomalib/models/image/uflow/lightning_model.py b/src/anomalib/models/image/uflow/lightning_model.py index 03ef56a2b0..ce88d4ae2e 100644 --- a/src/anomalib/models/image/uflow/lightning_model.py +++ b/src/anomalib/models/image/uflow/lightning_model.py @@ -13,13 +13,14 @@ from lightning.pytorch.core.optimizer import LightningOptimizer from lightning.pytorch.utilities.types import STEP_OUTPUT from torch.optim.lr_scheduler import LRScheduler -from torchvision.transforms.v2 import Compose, Normalize, Resize, Transform +from torchvision.transforms.v2 import Compose, Normalize, Resize from anomalib import LearningType from anomalib.data import Batch from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import PostProcessor +from anomalib.pre_processing import PreProcessor from .loss import UFlowLoss from .torch_model import UflowModel @@ -47,10 +48,32 @@ def __init__( affine_clamp: float = 2.0, affine_subnet_channels_ratio: float = 1.0, permute_soft: bool = False, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + """Uflow model. + + Args: + backbone (str): Backbone name. + flow_steps (int): Number of flow steps. + affine_clamp (float): Affine clamp. + affine_subnet_channels_ratio (float): Affine subnet channels ratio. + permute_soft (bool): Whether to use soft permutation. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. + post_processor (PostProcessor, optional): Post-processor for the model. + This is used to post-process the output data after it is passed to the model. + Defaults to ``None``. + evaluator (Evaluator, optional): Evaluator for the model. + This is used to evaluate the model. + Defaults to ``True``. + """ + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) + if self.input_size is None: + msg = "Input size is required for UFlow model." + raise ValueError(msg) self.backbone = backbone self.flow_steps = flow_steps @@ -58,15 +81,6 @@ def __init__( self.affine_subnet_channels_ratio = affine_subnet_channels_ratio self.permute_soft = permute_soft - self.loss = UFlowLoss() - - self.model: UflowModel - - def _setup(self) -> None: - if self.input_size is None: - msg = "Input size is required for UFlow model." - raise ValueError(msg) - self.model = UflowModel( input_size=self.input_size, backbone=self.backbone, @@ -75,18 +89,18 @@ def _setup(self) -> None: affine_subnet_channels_ratio=self.affine_subnet_channels_ratio, permute_soft=self.permute_soft, ) + self.loss = UFlowLoss() - def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: # noqa: ARG002 | unused arguments - """Training step.""" - z, ljd = self.model(batch.image) - loss = self.loss(z, ljd) - self.log_dict({"loss": loss}, on_step=True, on_epoch=False, prog_bar=False, logger=True) - return {"loss": loss} - - def validation_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: # noqa: ARG002 | unused arguments - """Validation step.""" - predictions = self.model(batch.image) - return batch.update(**predictions._asdict()) + @classmethod + def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> PreProcessor: + """Default pre-processor for UFlow.""" + if image_size is not None: + logger.warning("Image size is not used in UFlow. The input image size is determined by the model.") + transform = Compose([ + Resize((448, 448), antialias=True), + Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + return PreProcessor(transform=transform) def configure_optimizers(self) -> tuple[list[LightningOptimizer], list[LRScheduler]]: """Return optimizer and scheduler.""" @@ -106,6 +120,18 @@ def configure_optimizers(self) -> tuple[list[LightningOptimizer], list[LRSchedul ) return [optimizer], [scheduler] + def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: # noqa: ARG002 | unused arguments + """Training step.""" + z, ljd = self.model(batch.image) + loss = self.loss(z, ljd) + self.log_dict({"loss": loss}, on_step=True, on_epoch=False, prog_bar=False, logger=True) + return {"loss": loss} + + def validation_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT: # noqa: ARG002 | unused arguments + """Validation step.""" + predictions = self.model(batch.image) + return batch.update(**predictions._asdict()) + @property def trainer_arguments(self) -> dict[str, Any]: """Return EfficientAD trainer arguments.""" @@ -119,15 +145,3 @@ def learning_type(self) -> LearningType: LearningType: Learning type of the model. """ return LearningType.ONE_CLASS - - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform: - """Default transform for Padim.""" - if image_size is not None: - logger.warning("Image size is not used in UFlow. The input image size is determined by the model.") - return Compose( - [ - Resize((448, 448), antialias=True), - Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ], - ) diff --git a/src/anomalib/models/image/winclip/lightning_model.py b/src/anomalib/models/image/winclip/lightning_model.py index 2651c588e3..3ce35f46f4 100644 --- a/src/anomalib/models/image/winclip/lightning_model.py +++ b/src/anomalib/models/image/winclip/lightning_model.py @@ -13,7 +13,7 @@ import torch from torch.utils.data import DataLoader -from torchvision.transforms.v2 import Compose, InterpolationMode, Normalize, Resize, Transform +from torchvision.transforms.v2 import Compose, InterpolationMode, Normalize, Resize from anomalib import LearningType from anomalib.data import Batch @@ -21,6 +21,7 @@ from anomalib.metrics import Evaluator from anomalib.models.components import AnomalyModule from anomalib.post_processing import OneClassPostProcessor, PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import WinClipModel @@ -41,6 +42,9 @@ class WinClip(AnomalyModule): Defaults to ``(2, 3)``. few_shot_source (str | Path, optional): Path to a folder of reference images used for few-shot inference. Defaults to ``None``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ EXCLUDE_FROM_STATE_DICT = frozenset({"model.clip"}) @@ -51,10 +55,12 @@ def __init__( k_shot: int = 0, scales: tuple = (2, 3), few_shot_source: Path | str | None = None, + pre_processor: PreProcessor | bool = True, post_processor: PostProcessor | None = None, evaluator: Evaluator | bool = True, ) -> None: - super().__init__(post_processor=post_processor, evaluator=evaluator) + super().__init__(pre_processor=pre_processor, post_processor=post_processor, evaluator=evaluator) + self.model = WinClipModel(scales=scales, apply_transform=False) self.class_name = class_name self.k_shot = k_shot @@ -77,7 +83,10 @@ def _setup(self) -> None: if self.k_shot: if self.few_shot_source: logger.info("Loading reference images from %s", self.few_shot_source) - reference_dataset = PredictDataset(self.few_shot_source, transform=self.model.transform) + reference_dataset = PredictDataset( + self.few_shot_source, + transform=self.pre_processor.test_transform if self.pre_processor else None, + ) dataloader = DataLoader(reference_dataset, batch_size=1, shuffle=False) else: logger.info("Collecting reference images from training dataset") @@ -173,17 +182,17 @@ def load_state_dict(self, state_dict: OrderedDict[str, Any], strict: bool = True state_dict.update(restore_dict) return super().load_state_dict(state_dict, strict) - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform: - """Configure the default transforms used by the model.""" + @classmethod + def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> PreProcessor: + """Configure the default pre-processor used by the model.""" if image_size is not None: logger.warning("Image size is not used in WinCLIP. The input image size is determined by the model.") - return Compose( - [ - Resize((240, 240), antialias=True, interpolation=InterpolationMode.BICUBIC), - Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)), - ], - ) + + transform = Compose([ + Resize((240, 240), antialias=True, interpolation=InterpolationMode.BICUBIC), + Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)), + ]) + return PreProcessor(val_transform=transform, test_transform=transform) @staticmethod def default_post_processor() -> OneClassPostProcessor: diff --git a/src/anomalib/models/video/ai_vad/lightning_model.py b/src/anomalib/models/video/ai_vad/lightning_model.py index 8c36a689c7..9625f4b565 100644 --- a/src/anomalib/models/video/ai_vad/lightning_model.py +++ b/src/anomalib/models/video/ai_vad/lightning_model.py @@ -11,12 +11,12 @@ from typing import Any from lightning.pytorch.utilities.types import STEP_OUTPUT -from torchvision.transforms.v2 import Transform from anomalib import LearningType from anomalib.data import VideoBatch from anomalib.models.components import AnomalyModule, MemoryBankMixin from anomalib.post_processing.one_class import OneClassPostProcessor, PostProcessor +from anomalib.pre_processing import PreProcessor from .torch_model import AiVadModel @@ -59,6 +59,9 @@ class AiVad(MemoryBankMixin, AnomalyModule): Defaults to ``1``. n_neighbors_deep (int): Number of neighbors used in KNN density estimation for deep features. Defaults to ``1``. + pre_processor (PreProcessor, optional): Pre-processor for the model. + This is used to pre-process the input data before it is passed to the model. + Defaults to ``None``. """ def __init__( @@ -77,10 +80,10 @@ def __init__( n_components_velocity: int = 2, n_neighbors_pose: int = 1, n_neighbors_deep: int = 1, + pre_processor: PreProcessor | bool = True, **kwargs, ) -> None: - super().__init__(**kwargs) - + super().__init__(pre_processor=pre_processor, **kwargs) self.model = AiVadModel( box_score_thresh=box_score_thresh, persons_only=persons_only, @@ -165,11 +168,15 @@ def learning_type(self) -> LearningType: """ return LearningType.ONE_CLASS - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform | None: - """AI-VAD does not need a transform, as the region- and feature-extractors apply their own transforms.""" + @classmethod + def configure_pre_processor(cls, image_size: tuple[int, int] | None = None) -> PreProcessor: + """Configure the pre-processor for AI-VAD. + + AI-VAD does not need a pre-processor or transforms, as the region- and + feature-extractors apply their own transforms. + """ del image_size - return None + return PreProcessor() # A pre-processor with no transforms. @staticmethod def default_post_processor() -> PostProcessor: diff --git a/src/anomalib/pre_processing/__init__.py b/src/anomalib/pre_processing/__init__.py new file mode 100644 index 0000000000..d70565f882 --- /dev/null +++ b/src/anomalib/pre_processing/__init__.py @@ -0,0 +1,8 @@ +"""Anomalib pre-processing module.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from .pre_processing import PreProcessor + +__all__ = ["PreProcessor"] diff --git a/src/anomalib/pre_processing/pre_processing.py b/src/anomalib/pre_processing/pre_processing.py new file mode 100644 index 0000000000..27cffc7605 --- /dev/null +++ b/src/anomalib/pre_processing/pre_processing.py @@ -0,0 +1,177 @@ +"""Anomalib pre-processing module.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from typing import TYPE_CHECKING + +import torch +from lightning import Callback, LightningModule, Trainer +from lightning.pytorch.trainer.states import TrainerFn +from torch import nn +from torch.utils.data import DataLoader +from torchvision.transforms.v2 import Transform + +from .utils.transform import ( + get_dataloaders_transforms, + get_exportable_transform, + set_dataloaders_transforms, + set_datamodule_stage_transform, +) + +if TYPE_CHECKING: + from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS + + from anomalib.data import AnomalibDataModule + + +class PreProcessor(nn.Module, Callback): + """Anomalib pre-processor. + + This class serves as both a PyTorch module and a Lightning callback, handling + the application of transforms to data batches during different stages of + training, validation, testing, and prediction. + + Args: + train_transform (Transform | None): Transform to apply during training. + val_transform (Transform | None): Transform to apply during validation. + test_transform (Transform | None): Transform to apply during testing. + transform (Transform | None): General transform to apply if stage-specific + transforms are not provided. + + Raises: + ValueError: If both `transform` and any of the stage-specific transforms + are provided simultaneously. + + Notes: + If only `transform` is provided, it will be used for all stages (train, val, test). + + Priority of transforms: + 1. Explicitly set PreProcessor transforms (highest priority) + 2. Datamodule transforms (if PreProcessor has no transforms) + 3. Dataloader transforms (if neither PreProcessor nor datamodule have transforms) + 4. Default transforms (lowest priority) + + Examples: + >>> from torchvision.transforms.v2 import Compose, Resize, ToTensor + >>> from anomalib.pre_processing import PreProcessor + + >>> # Define transforms + >>> train_transform = Compose([Resize((224, 224)), ToTensor()]) + >>> val_transform = Compose([Resize((256, 256)), CenterCrop((224, 224)), ToTensor()]) + + >>> # Create PreProcessor with stage-specific transforms + >>> pre_processor = PreProcessor( + ... train_transform=train_transform, + ... val_transform=val_transform + ... ) + + >>> # Create PreProcessor with a single transform for all stages + >>> common_transform = Compose([Resize((224, 224)), ToTensor()]) + >>> pre_processor_common = PreProcessor(transform=common_transform) + + >>> # Use in a Lightning module + >>> class MyModel(LightningModule): + ... def __init__(self): + ... super().__init__() + ... self.pre_processor = PreProcessor(...) + ... + ... def configure_callbacks(self): + ... return [self.pre_processor] + ... + ... def training_step(self, batch, batch_idx): + ... # The pre_processor will automatically apply the correct transform + ... processed_batch = self.pre_processor(batch) + ... # Rest of the training step + """ + + def __init__( + self, + train_transform: Transform | None = None, + val_transform: Transform | None = None, + test_transform: Transform | None = None, + transform: Transform | None = None, + ) -> None: + super().__init__() + + if transform and any([train_transform, val_transform, test_transform]): + msg = ( + "`transforms` cannot be used together with `train_transform`, `val_transform`, `test_transform`.\n" + "If you want to apply the same transform to the training, validation and test data, " + "use only `transforms`. \n" + "Otherwise, specify transforms for training, validation and test individually." + ) + raise ValueError(msg) + + self.train_transform = train_transform or transform + self.val_transform = val_transform or transform + self.test_transform = test_transform or transform + self.predict_transform = self.test_transform + self.export_transform = get_exportable_transform(self.test_transform) + + def setup_datamodule_transforms(self, datamodule: "AnomalibDataModule") -> None: + """Set up datamodule transforms.""" + # If PreProcessor has transforms, propagate them to datamodule + if any([self.train_transform, self.val_transform, self.test_transform]): + transforms = { + "fit": self.train_transform, + "val": self.val_transform, + "test": self.test_transform, + "predict": self.predict_transform, + } + + for stage, transform in transforms.items(): + if transform is not None: + set_datamodule_stage_transform(datamodule, transform, stage) + + def setup_dataloader_transforms(self, dataloaders: "EVAL_DATALOADERS | TRAIN_DATALOADERS") -> None: + """Set up dataloader transforms.""" + if isinstance(dataloaders, DataLoader): + dataloaders = [dataloaders] + + # If PreProcessor has transforms, propagate them to dataloaders + if any([self.train_transform, self.val_transform, self.test_transform]): + transforms = { + "train": self.train_transform, + "val": self.val_transform, + "test": self.test_transform, + } + set_dataloaders_transforms(dataloaders, transforms) + return + + # Try to get transforms from dataloaders + if dataloaders: + dataloaders_transforms = get_dataloaders_transforms(dataloaders) + if dataloaders_transforms: + self.train_transform = dataloaders_transforms.get("train") + self.val_transform = dataloaders_transforms.get("val") + self.test_transform = dataloaders_transforms.get("test") + self.predict_transform = self.test_transform + self.export_transform = get_exportable_transform(self.test_transform) + + def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None: + """Configure transforms at the start of each stage. + + Args: + trainer: The Lightning trainer. + pl_module: The Lightning module. + stage: The stage (e.g., 'fit', 'validate', 'test', 'predict'). + """ + stage = TrainerFn(stage).value # Ensure stage is str + + if hasattr(trainer, "datamodule"): + self.setup_datamodule_transforms(datamodule=trainer.datamodule) + elif hasattr(trainer, f"{stage}_dataloaders"): + dataloaders = getattr(trainer, f"{stage}_dataloaders") + self.setup_dataloader_transforms(dataloaders=dataloaders) + + super().setup(trainer, pl_module, stage) + + def forward(self, batch: torch.Tensor) -> torch.Tensor: + """Apply transforms to the batch of tensors for inference. + + This forward-pass is only used after the model is exported. + Within the Lightning training/validation/testing loops, the transforms are applied + in the `on_*_batch_start` methods. + """ + return self.export_transform(batch) if self.export_transform else batch diff --git a/src/anomalib/pre_processing/utils/__init__.py b/src/anomalib/pre_processing/utils/__init__.py new file mode 100644 index 0000000000..8361223189 --- /dev/null +++ b/src/anomalib/pre_processing/utils/__init__.py @@ -0,0 +1,4 @@ +"""Utility functions for pre-processing.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 diff --git a/src/anomalib/pre_processing/utils/transform.py b/src/anomalib/pre_processing/utils/transform.py new file mode 100644 index 0000000000..37eb1e9dd1 --- /dev/null +++ b/src/anomalib/pre_processing/utils/transform.py @@ -0,0 +1,150 @@ +"""Utility functions for transforms.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence + +from torch.utils.data import DataLoader +from torchvision.transforms.v2 import CenterCrop, Compose, Resize, Transform + +from anomalib.data import AnomalibDataModule +from anomalib.data.transforms import ExportableCenterCrop + + +def get_dataloaders_transforms(dataloaders: Sequence[DataLoader]) -> dict[str, Transform]: + """Get transforms from dataloaders. + + Args: + dataloaders: The dataloaders to get transforms from. + + Returns: + Dictionary mapping stages to their transforms. + """ + transforms: dict[str, Transform] = {} + stage_lookup = { + "fit": "train", + "validate": "val", + "test": "test", + "predict": "test", + } + + for dataloader in dataloaders: + if not hasattr(dataloader, "dataset") or not hasattr(dataloader.dataset, "transform"): + continue + + for stage in stage_lookup: + if hasattr(dataloader, f"{stage}_dataloader"): + transforms[stage_lookup[stage]] = dataloader.dataset.transform + + return transforms + + +def set_dataloaders_transforms(dataloaders: Sequence[DataLoader], transforms: dict[str, Transform | None]) -> None: + """Set transforms to dataloaders. + + Args: + dataloaders: The dataloaders to propagate transforms to. + transforms: Dictionary mapping stages to their transforms. + """ + stage_mapping = { + "fit": "train", + "validate": "val", + "test": "test", + "predict": "test", # predict uses test transform + } + + for loader in dataloaders: + if not hasattr(loader, "dataset"): + continue + + for stage in stage_mapping: + if hasattr(loader, f"{stage}_dataloader"): + transform = transforms.get(stage_mapping[stage]) + if transform is not None: + set_dataloader_transform([loader], transform) + + +def set_dataloader_transform(dataloader: DataLoader | Sequence[DataLoader], transform: Transform) -> None: + """Set a transform for a dataloader or list of dataloaders. + + Args: + dataloader: The dataloader(s) to set the transform for. + transform: The transform to set. + """ + if isinstance(dataloader, DataLoader): + if hasattr(dataloader.dataset, "transform"): + dataloader.dataset.transform = transform + elif isinstance(dataloader, Sequence): + for dl in dataloader: + set_dataloader_transform(dl, transform) + else: + msg = f"Unsupported dataloader type: {type(dataloader)}" + raise TypeError(msg) + + +def set_datamodule_stage_transform(datamodule: AnomalibDataModule, transform: Transform, stage: str) -> None: + """Set a transform for a specific stage in a AnomalibDataModule. + + Args: + datamodule: The AnomalibDataModule to set the transform for. + transform: The transform to set. + stage: The stage to set the transform for. + + Note: + The stage parameter maps to dataset attributes as follows: + - 'fit' -> 'train_data' + - 'validate' -> 'val_data' + - 'test' -> 'test_data' + - 'predict' -> 'test_data' + """ + stage_datasets = { + "fit": "train_data", + "validate": "val_data", + "test": "test_data", + "predict": "test_data", + } + + dataset_attr = stage_datasets.get(stage) + if dataset_attr and hasattr(datamodule, dataset_attr): + dataset = getattr(datamodule, dataset_attr) + if hasattr(dataset, "transform"): + dataset.transform = transform + + +def get_exportable_transform(transform: Transform | None) -> Transform | None: + """Get exportable transform. + + Some transforms are not supported by ONNX/OpenVINO, so we need to replace them with exportable versions. + """ + if transform is None: + return None + transform = disable_antialiasing(transform) + return convert_center_crop_transform(transform) + + +def disable_antialiasing(transform: Transform) -> Transform: + """Disable antialiasing in Resize transforms. + + Resizing with antialiasing is not supported by ONNX, so we need to disable it. + """ + if isinstance(transform, Resize): + transform.antialias = False + if isinstance(transform, Compose): + for tr in transform.transforms: + disable_antialiasing(tr) + return transform + + +def convert_center_crop_transform(transform: Transform) -> Transform: + """Convert CenterCrop to ExportableCenterCrop. + + Torchvision's CenterCrop is not supported by ONNX, so we need to replace it with our own ExportableCenterCrop. + """ + if isinstance(transform, CenterCrop): + transform = ExportableCenterCrop(size=transform.size) + if isinstance(transform, Compose): + for index in range(len(transform.transforms)): + tr = transform.transforms[index] + transform.transforms[index] = convert_center_crop_transform(tr) + return transform diff --git a/tests/integration/model/test_models.py b/tests/integration/model/test_models.py index 4fc266d190..2ffd2188f4 100644 --- a/tests/integration/model/test_models.py +++ b/tests/integration/model/test_models.py @@ -187,12 +187,9 @@ def _get_objects( extra_args = {} if model_name in {"rkde", "dfkde"}: extra_args["n_pca_components"] = 2 + if model_name == "ai_vad": pytest.skip("Revisit AI-VAD test") - - # select dataset - elif model_name == "win_clip": - dataset = MVTec(root=dataset_path / "mvtec", category="dummy", image_size=240, task=task_type) else: # EfficientAd requires that the batch size be lesser than the number of images in the dataset. # This is so that the LR step size is not 0. diff --git a/tests/integration/tools/upgrade/expected_draem_v1.yaml b/tests/integration/tools/upgrade/expected_draem_v1.yaml index f59a21d5e9..882e27b74e 100644 --- a/tests/integration/tools/upgrade/expected_draem_v1.yaml +++ b/tests/integration/tools/upgrade/expected_draem_v1.yaml @@ -3,16 +3,10 @@ data: init_args: root: ./datasets/MVTec category: bottle - image_size: - - 256 - - 256 train_batch_size: 72 eval_batch_size: 32 num_workers: 8 task: segmentation - transform: null - train_transform: null - eval_transform: null test_split_mode: from_dir test_split_ratio: 0.2 val_split_mode: same_as_test @@ -27,6 +21,7 @@ model: beta: - 0.1 - 1.0 + pre_processor: true post_processor: null evaluator: true normalization: diff --git a/tests/unit/data/datamodule/depth/test_folder_3d.py b/tests/unit/data/datamodule/depth/test_folder_3d.py index 6ed01bfff5..9ebf82e3f2 100644 --- a/tests/unit/data/datamodule/depth/test_folder_3d.py +++ b/tests/unit/data/datamodule/depth/test_folder_3d.py @@ -29,7 +29,6 @@ def datamodule(dataset_path: Path, task_type: TaskType) -> Folder3D: normal_depth_dir="train/good/xyz", abnormal_depth_dir="test/bad/xyz", normal_test_depth_dir="test/good/xyz", - image_size=256, train_batch_size=4, eval_batch_size=4, num_workers=0, diff --git a/tests/unit/data/datamodule/depth/test_mvtec_3d.py b/tests/unit/data/datamodule/depth/test_mvtec_3d.py index 70966b7774..6a94f1b279 100644 --- a/tests/unit/data/datamodule/depth/test_mvtec_3d.py +++ b/tests/unit/data/datamodule/depth/test_mvtec_3d.py @@ -23,7 +23,6 @@ def datamodule(dataset_path: Path, task_type: TaskType) -> MVTec3D: root=dataset_path / "mvtec_3d", category="dummy", task=task_type, - image_size=256, train_batch_size=4, eval_batch_size=4, num_workers=0, diff --git a/tests/unit/data/datamodule/image/test_btech.py b/tests/unit/data/datamodule/image/test_btech.py index cf7b207e1d..2f483da7c8 100644 --- a/tests/unit/data/datamodule/image/test_btech.py +++ b/tests/unit/data/datamodule/image/test_btech.py @@ -23,7 +23,6 @@ def datamodule(dataset_path: Path, task_type: TaskType) -> BTech: root=dataset_path / "btech", category="dummy", task=task_type, - image_size=256, train_batch_size=4, eval_batch_size=4, ) diff --git a/tests/unit/data/datamodule/image/test_kolektor.py b/tests/unit/data/datamodule/image/test_kolektor.py index 703c3927a3..7fc061c09d 100644 --- a/tests/unit/data/datamodule/image/test_kolektor.py +++ b/tests/unit/data/datamodule/image/test_kolektor.py @@ -22,7 +22,6 @@ def datamodule(dataset_path: Path, task_type: TaskType) -> Kolektor: _datamodule = Kolektor( root=dataset_path / "kolektor", task=task_type, - image_size=256, train_batch_size=4, eval_batch_size=4, ) diff --git a/tests/unit/data/datamodule/image/test_visa.py b/tests/unit/data/datamodule/image/test_visa.py index 0c663a6e54..8b173f38cc 100644 --- a/tests/unit/data/datamodule/image/test_visa.py +++ b/tests/unit/data/datamodule/image/test_visa.py @@ -22,7 +22,6 @@ def datamodule(dataset_path: Path, task_type: TaskType) -> Visa: _datamodule = Visa( root=dataset_path, category="dummy", - image_size=256, train_batch_size=4, eval_batch_size=4, num_workers=0, diff --git a/tests/unit/data/datamodule/video/test_avenue.py b/tests/unit/data/datamodule/video/test_avenue.py index 42365d059f..5069b93def 100644 --- a/tests/unit/data/datamodule/video/test_avenue.py +++ b/tests/unit/data/datamodule/video/test_avenue.py @@ -29,7 +29,6 @@ def datamodule(dataset_path: Path, task_type: TaskType, clip_length_in_frames: i root=dataset_path / "avenue", gt_dir=dataset_path / "avenue" / "ground_truth_demo", clip_length_in_frames=clip_length_in_frames, - image_size=256, task=task_type, num_workers=0, train_batch_size=4, diff --git a/tests/unit/data/datamodule/video/test_shanghaitech.py b/tests/unit/data/datamodule/video/test_shanghaitech.py index fda0d1a84d..4e96cfbaa7 100644 --- a/tests/unit/data/datamodule/video/test_shanghaitech.py +++ b/tests/unit/data/datamodule/video/test_shanghaitech.py @@ -29,7 +29,6 @@ def datamodule(dataset_path: Path, task_type: TaskType, clip_length_in_frames: i root=dataset_path / "shanghaitech", scene=1, clip_length_in_frames=clip_length_in_frames, - image_size=(256, 256), train_batch_size=4, eval_batch_size=4, num_workers=0, diff --git a/tests/unit/data/datamodule/video/test_ucsdped.py b/tests/unit/data/datamodule/video/test_ucsdped.py index 1148e9313a..669d72278a 100644 --- a/tests/unit/data/datamodule/video/test_ucsdped.py +++ b/tests/unit/data/datamodule/video/test_ucsdped.py @@ -30,7 +30,6 @@ def datamodule(dataset_path: Path, task_type: TaskType, clip_length_in_frames: i category="dummy", clip_length_in_frames=clip_length_in_frames, task=task_type, - image_size=256, train_batch_size=4, eval_batch_size=4, num_workers=0, diff --git a/tests/unit/engine/test_engine.py b/tests/unit/engine/test_engine.py index 1fdb7532a4..c927733595 100644 --- a/tests/unit/engine/test_engine.py +++ b/tests/unit/engine/test_engine.py @@ -85,10 +85,6 @@ def fxt_full_config_path(tmp_path: Path) -> Path: train_batch_size: 32 eval_batch_size: 32 num_workers: 8 - image_size: null - transform: null - train_transform: null - eval_transform: null test_split_mode: FROM_DIR test_split_ratio: 0.2 val_split_mode: SAME_AS_TEST diff --git a/tests/unit/engine/test_setup_transform.py b/tests/unit/engine/test_setup_transform.py deleted file mode 100644 index ebb60f81c0..0000000000 --- a/tests/unit/engine/test_setup_transform.py +++ /dev/null @@ -1,260 +0,0 @@ -"""Tests for the Anomalib Engine.""" - -# Copyright (C) 2024 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -import tempfile -from collections.abc import Generator -from pathlib import Path - -import pytest -import torch -from torch.utils.data import DataLoader -from torchvision.transforms.v2 import Resize, Transform - -from anomalib import LearningType, TaskType -from anomalib.data import AnomalibDataModule, AnomalibDataset, InferenceBatch -from anomalib.engine import Engine -from anomalib.models import AnomalyModule -from anomalib.post_processing import PostProcessor - - -class DummyDataset(AnomalibDataset): - """Dummy dataset for testing the setup_transform method.""" - - def __init__(self, transform: Transform = None) -> None: - super().__init__(TaskType.CLASSIFICATION, transform=transform) - self.image = torch.rand(3, 10, 10) - self._samples = None - - def _setup(self, _stage: str | None = None) -> None: - self._samples = None - - def __len__(self) -> int: - """Return the length of the dataset.""" - return 1 - - -class DummyPostProcessor(PostProcessor): - """Dummy post-processor for testing the setup_transform method.""" - - @staticmethod - def forward(batch: InferenceBatch) -> InferenceBatch: - """Return the batch unmodified.""" - return batch - - -class DummyModel(AnomalyModule): - """Dummy model for testing the setup_transform method.""" - - def __init__(self) -> None: - super().__init__() - self.model = torch.nn.Linear(10, 10) - - @staticmethod - def configure_transforms(image_size: tuple[int, int] | None = None) -> Transform: - """Return a Resize transform.""" - if image_size is None: - image_size = (256, 256) - return Resize(image_size) - - @staticmethod - def trainer_arguments() -> dict: - """Return an empty dictionary.""" - return {} - - @staticmethod - def learning_type() -> LearningType: - """Return the learning type.""" - return LearningType.ZERO_SHOT - - @staticmethod - def default_post_processor() -> PostProcessor: - """Return a dummy post-processor.""" - return DummyPostProcessor() - - -class DummyDataModule(AnomalibDataModule): - """Dummy datamodule for testing the setup_transform method.""" - - def __init__( - self, - transform: Transform | None = None, - train_transform: Transform | None = None, - eval_transform: Transform | None = None, - image_size: tuple[int, int] | None = None, - ) -> None: - super().__init__( - train_batch_size=1, - eval_batch_size=1, - num_workers=0, - val_split_mode="from_test", - val_split_ratio=0.5, - image_size=image_size, - transform=transform, - train_transform=train_transform, - eval_transform=eval_transform, - ) - - def _create_val_split(self) -> None: - pass - - def _create_test_split(self) -> None: - pass - - def _setup(self, _stage: str | None = None) -> None: - self.train_data = DummyDataset(transform=self.train_transform) - self.val_data = DummyDataset(transform=self.eval_transform) - self.test_data = DummyDataset(transform=self.eval_transform) - - -@pytest.fixture() -def checkpoint_path() -> Generator: - """Fixture to create a temporary checkpoint file that stores a Resize transform.""" - # Create a temporary file - transform = Resize((50, 50)) - with tempfile.TemporaryDirectory() as temp_dir: - file_path = Path(temp_dir) / "model.ckpt" - checkpoint = {"transform": transform} - torch.save(checkpoint, file_path) - - yield file_path - - -class TestSetupTransform: - """Tests for the `_setup_transform` method of the Anomalib Engine.""" - - # test update single dataloader - @staticmethod - def test_single_dataloader_default_transform() -> None: - """Tests if the default model transform is used when no transform is passed to the dataloader.""" - dataset = DummyDataset() - dataloader = DataLoader(dataset, batch_size=1) - model = DummyModel() - # before the setup_transform is called, the dataset should not have a transform - assert dataset.transform is None - Engine._setup_transform(model, dataloaders=dataloader) # noqa: SLF001 - # after the setup_transform is called, the dataset should have the default transform from the model - assert dataset.transform is not None - - # test update multiple dataloaders - @staticmethod - def test_multiple_dataloaders_default_transform() -> None: - """Tests if the default model transform is used when no transform is passed to the dataloader.""" - dataset = DummyDataset() - dataloader = DataLoader(dataset, batch_size=1) - model = DummyModel() - # before the setup_transform is called, the dataset should not have a transform - assert dataset.transform is None - Engine._setup_transform(model, dataloaders=[dataloader, dataloader]) # noqa: SLF001 - # after the setup_transform is called, the dataset should have the default transform from the model - assert dataset.transform is not None - - @staticmethod - def test_single_dataloader_custom_transform() -> None: - """Tests if the user-specified transform is used when passed to the dataloader.""" - transform = Transform() - dataset = DummyDataset(transform=transform) - dataloader = DataLoader(dataset, batch_size=1) - model = DummyModel() - # before the setup_transform is called, the dataset should have the custom transform - assert dataset.transform == transform - Engine._setup_transform(model, dataloaders=dataloader) # noqa: SLF001 - # after the setup_transform is called, the model should have the custom transform - assert model.transform == transform - - # test if the user-specified transform is used when passed to the datamodule - @staticmethod - def test_custom_transform() -> None: - """Tests if the user-specified transform is used when passed to the datamodule.""" - transform = Transform() - datamodule = DummyDataModule(transform=transform) - model = DummyModel() - # assert that the datamodule uses the custom transform before and after setup_transform is called - assert datamodule.train_transform == transform - assert datamodule.eval_transform == transform - Engine._setup_transform(model, datamodule=datamodule) # noqa: SLF001 - assert datamodule.train_transform == transform - assert datamodule.eval_transform == transform - assert model.transform == transform - - # test if the user-specified transform is used when passed to the datamodule - @staticmethod - def test_custom_train_transform() -> None: - """Tests if the user-specified transform is used when passed to the datamodule as train_transform.""" - model = DummyModel() - transform = Transform() - datamodule = DummyDataModule(train_transform=transform) - # before calling setup, train_transform should be the custom transform and eval_transform should be None - assert datamodule.train_transform == transform - assert datamodule.eval_transform is None - Engine._setup_transform(model, datamodule=datamodule) # noqa: SLF001 - # after calling setup, train_transform should be the custom transform and eval_transform should be the default - assert datamodule.train_transform == transform - assert datamodule.eval_transform is None - assert model.transform != transform - assert model.transform is not None - - # test if the user-specified transform is used when passed to the datamodule - @staticmethod - def test_custom_eval_transform() -> None: - """Tests if the user-specified transform is used when passed to the datamodule as eval_transform.""" - model = DummyModel() - transform = Transform() - datamodule = DummyDataModule(eval_transform=transform) - # before calling setup, train_transform should be the custom transform and eval_transform should be None - assert datamodule.train_transform is None - assert datamodule.eval_transform == transform - Engine._setup_transform(model, datamodule=datamodule) # noqa: SLF001 - # after calling setup, train_transform should be the custom transform and eval_transform should be the default - assert datamodule.train_transform is None - assert datamodule.eval_transform == transform - assert model.transform == transform - - # test update datamodule - @staticmethod - def test_datamodule_default_transform() -> None: - """Tests if the default model transform is used when no transform is passed to the datamodule.""" - datamodule = DummyDataModule() - model = DummyModel() - # assert that the datamodule has a transform after the setup_transform is called - Engine._setup_transform(model, datamodule=datamodule) # noqa: SLF001 - assert isinstance(model.transform, Transform) - - # test if image size is taken from datamodule - @staticmethod - def test_datamodule_image_size() -> None: - """Tests if the image size that is passed to the datamodule overwrites the default size from the model.""" - datamodule = DummyDataModule(image_size=(100, 100)) - model = DummyModel() - # assert that the datamodule has a transform after the setup_transform is called - Engine._setup_transform(model, datamodule=datamodule) # noqa: SLF001 - assert isinstance(model.transform, Resize) - assert model.transform.size == [100, 100] - - @staticmethod - def test_transform_from_checkpoint(checkpoint_path: Path) -> None: - """Tests if the transform from the checkpoint is used.""" - model = DummyModel() - Engine._setup_transform(model, ckpt_path=checkpoint_path) # noqa: SLF001 - assert isinstance(model.transform, Resize) - assert model.transform.size == [50, 50] - - @staticmethod - def test_precendence_datamodule(checkpoint_path: Path) -> None: - """Tests if transform from the datamodule goes first if both checkpoint and datamodule are provided.""" - transform = Transform() - datamodule = DummyDataModule(transform=transform) - model = DummyModel() - Engine._setup_transform(model, ckpt_path=checkpoint_path, datamodule=datamodule) # noqa: SLF001 - assert model.transform == transform - - @staticmethod - def test_transform_already_assigned() -> None: - """Tests if the transform from the model is used when the model already has a transform assigned.""" - transform = Transform() - model = DummyModel() - model.set_transform(transform) - datamodule = DummyDataModule() - Engine._setup_transform(model, datamodule=datamodule) # noqa: SLF001 - assert model.transform == transform diff --git a/tests/unit/pre_processing/test_pre_processing.py b/tests/unit/pre_processing/test_pre_processing.py new file mode 100644 index 0000000000..36394d54a3 --- /dev/null +++ b/tests/unit/pre_processing/test_pre_processing.py @@ -0,0 +1,127 @@ +"""Test the PreProcessor class.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import MagicMock + +import pytest +import torch +from torch.utils.data import DataLoader +from torchvision.transforms.v2 import Compose, Resize, ToDtype, ToImage +from torchvision.tv_tensors import Image, Mask + +from anomalib.data import ImageBatch +from anomalib.pre_processing import PreProcessor + + +class TestPreProcessor: + """Test the PreProcessor class.""" + + @pytest.fixture(autouse=True) + def setup(self) -> None: + """Set up test fixtures for each test method.""" + image = Image(torch.rand(3, 256, 256)) + gt_mask = Mask(torch.zeros(256, 256)) + self.dummy_batch = ImageBatch(image=image, gt_mask=gt_mask) + self.common_transform = Compose([Resize((224, 224)), ToImage(), ToDtype(torch.float32, scale=True)]) + + def test_init(self) -> None: + """Test the initialization of the PreProcessor class.""" + # Test with stage-specific transforms + train_transform = Compose([Resize((224, 224)), ToImage(), ToDtype(torch.float32, scale=True)]) + val_transform = Compose([Resize((256, 256)), ToImage(), ToDtype(torch.float32, scale=True)]) + pre_processor = PreProcessor(train_transform=train_transform, val_transform=val_transform) + assert pre_processor.train_transform == train_transform + assert pre_processor.val_transform == val_transform + assert pre_processor.test_transform is None + + # Test with single transform for all stages + pre_processor = PreProcessor(transform=self.common_transform) + assert pre_processor.train_transform == self.common_transform + assert pre_processor.val_transform == self.common_transform + assert pre_processor.test_transform == self.common_transform + + # Test error case: both transform and stage-specific transform + with pytest.raises(ValueError, match="`transforms` cannot be used together with"): + PreProcessor(transform=self.common_transform, train_transform=train_transform) + + def test_forward(self) -> None: + """Test the forward method of the PreProcessor class.""" + pre_processor = PreProcessor(transform=self.common_transform) + processed_batch = pre_processor(self.dummy_batch.image) + assert isinstance(processed_batch, torch.Tensor) + assert processed_batch.shape == (1, 3, 224, 224) + + def test_no_transform(self) -> None: + """Test no transform.""" + pre_processor = PreProcessor() + processed_batch = pre_processor(self.dummy_batch.image) + assert isinstance(processed_batch, torch.Tensor) + assert processed_batch.shape == (1, 3, 256, 256) + + @staticmethod + def test_different_stage_transforms() -> None: + """Test different stage transforms.""" + train_transform = Compose([Resize((224, 224)), ToImage(), ToDtype(torch.float32, scale=True)]) + val_transform = Compose([Resize((256, 256)), ToImage(), ToDtype(torch.float32, scale=True)]) + test_transform = Compose([Resize((288, 288)), ToImage(), ToDtype(torch.float32, scale=True)]) + + pre_processor = PreProcessor( + train_transform=train_transform, + val_transform=val_transform, + test_transform=test_transform, + ) + + # Test train transform + test_batch = ImageBatch(image=Image(torch.rand(3, 256, 256)), gt_mask=Mask(torch.zeros(256, 256))) + processed_batch = pre_processor.train_transform(test_batch.image) + assert isinstance(processed_batch, torch.Tensor) + assert processed_batch.shape == (1, 3, 224, 224) + + # Test validation transform + test_batch = ImageBatch(image=Image(torch.rand(3, 256, 256)), gt_mask=Mask(torch.zeros(256, 256))) + processed_batch = pre_processor.val_transform(test_batch.image) + assert isinstance(processed_batch, torch.Tensor) + assert processed_batch.shape == (1, 3, 256, 256) + + # Test test transform + test_batch = ImageBatch(image=Image(torch.rand(3, 256, 256)), gt_mask=Mask(torch.zeros(256, 256))) + processed_batch = pre_processor.test_transform(test_batch.image) + assert isinstance(processed_batch, torch.Tensor) + assert processed_batch.shape == (1, 3, 288, 288) + + def test_setup_transforms_from_dataloaders(self) -> None: + """Test setup method when transforms are obtained from dataloaders.""" + # Mock dataloader with dataset having a transform + dataloader = MagicMock() + dataloader.dataset.transform = self.common_transform + + pre_processor = PreProcessor() + pre_processor.setup_dataloader_transforms(dataloaders=[dataloader]) + + assert pre_processor.train_transform == self.common_transform + assert pre_processor.val_transform == self.common_transform + assert pre_processor.test_transform == self.common_transform + + def test_setup_transforms_priority(self) -> None: + """Test setup method prioritizes PreProcessor transforms over datamodule/dataloaders.""" + # Mock datamodule + datamodule = MagicMock() + datamodule.train_transform = Compose([Resize((128, 128)), ToImage(), ToDtype(torch.float32, scale=True)]) + datamodule.eval_transform = Compose([Resize((128, 128)), ToImage(), ToDtype(torch.float32, scale=True)]) + + # Mock dataloader + dataset_mock = MagicMock() + dataset_mock.transform = Compose([Resize((64, 64)), ToImage(), ToDtype(torch.float32, scale=True)]) + dataloader = MagicMock(spec=DataLoader) + dataloader.dataset = dataset_mock + + # Initialize PreProcessor with a custom transform + pre_processor = PreProcessor(transform=self.common_transform) + pre_processor.setup_datamodule_transforms(datamodule=datamodule) + + # Ensure PreProcessor's own transform is used + assert pre_processor.train_transform == self.common_transform + assert pre_processor.val_transform == self.common_transform + assert pre_processor.test_transform == self.common_transform diff --git a/tests/unit/pre_processing/utils/test_transform.py b/tests/unit/pre_processing/utils/test_transform.py new file mode 100644 index 0000000000..6974bcdbc8 --- /dev/null +++ b/tests/unit/pre_processing/utils/test_transform.py @@ -0,0 +1,103 @@ +"""Test the pre-processing transforms utils.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from torch.utils.data import DataLoader, TensorDataset +from torchvision.transforms.v2 import CenterCrop, Compose, Resize, ToTensor + +from anomalib.data.transforms import ExportableCenterCrop +from anomalib.pre_processing.utils.transform import ( + convert_center_crop_transform, + disable_antialiasing, + get_exportable_transform, + set_dataloader_transform, +) + + +def test_set_dataloader_transform() -> None: + """Test the set_dataloader_transform function.""" + + # Test with single DataLoader + class TransformableDataset(TensorDataset): + def __init__(self, *tensors) -> None: + super().__init__(*tensors) + self.transform = None + + dataset = TransformableDataset(torch.randn(10, 3, 224, 224)) + dataloader = DataLoader(dataset) + transform = ToTensor() + set_dataloader_transform(dataloader, transform) + assert dataloader.dataset.transform == transform + + # Test with sequence of DataLoaders + dataloaders = [DataLoader(TransformableDataset(torch.randn(10, 3, 224, 224))) for _ in range(3)] + set_dataloader_transform(dataloaders, transform) + for dl in dataloaders: + assert dl.dataset.transform == transform + + # Test with unsupported type + with pytest.raises(TypeError): + set_dataloader_transform({"key": "value"}, transform) + + +def test_get_exportable_transform() -> None: + """Test the get_exportable_transform function.""" + # Test with None transform + assert get_exportable_transform(None) is None + + # Test with Resize transform + resize = Resize((224, 224), antialias=True) + exportable_resize = get_exportable_transform(resize) + assert isinstance(exportable_resize, Resize) + assert not exportable_resize.antialias + + # Test with CenterCrop transform + center_crop = CenterCrop((224, 224)) + exportable_center_crop = get_exportable_transform(center_crop) + assert isinstance(exportable_center_crop, ExportableCenterCrop) + + # Test with Compose transform + compose = Compose([Resize((224, 224), antialias=True), CenterCrop((200, 200))]) + exportable_compose = get_exportable_transform(compose) + assert isinstance(exportable_compose, Compose) + assert isinstance(exportable_compose.transforms[0], Resize) + assert not exportable_compose.transforms[0].antialias + assert isinstance(exportable_compose.transforms[1], ExportableCenterCrop) + + +def test_disable_antialiasing() -> None: + """Test the disable_antialiasing function.""" + # Test with Resize transform + resize = Resize((224, 224), antialias=True) + disabled_resize = disable_antialiasing(resize) + assert not disabled_resize.antialias + + # Test with Compose transform + compose = Compose([Resize((224, 224), antialias=True), ToTensor()]) + disabled_compose = disable_antialiasing(compose) + assert not disabled_compose.transforms[0].antialias + + # Test with non-Resize transform + to_tensor = ToTensor() + assert disable_antialiasing(to_tensor) == to_tensor + + +def test_convert_centercrop() -> None: + """Test the convert_centercrop function.""" + # Test with CenterCrop transform + center_crop = CenterCrop((224, 224)) + converted_crop = convert_center_crop_transform(center_crop) + assert isinstance(converted_crop, ExportableCenterCrop) + assert converted_crop.size == list(center_crop.size) + + # Test with Compose transform + compose = Compose([Resize((256, 256)), CenterCrop((224, 224))]) + converted_compose = convert_center_crop_transform(compose) + assert isinstance(converted_compose.transforms[1], ExportableCenterCrop) + + # Test with non-CenterCrop transform + resize = Resize((224, 224)) + assert convert_center_crop_transform(resize) == resize diff --git a/tools/upgrade/config.py b/tools/upgrade/config.py index 71bf17a4b5..5f1f3278e1 100644 --- a/tools/upgrade/config.py +++ b/tools/upgrade/config.py @@ -27,7 +27,6 @@ import yaml from anomalib.models import convert_snake_to_pascal_case -from anomalib.utils.config import to_tuple def get_class_signature(module_path: str, class_name: str) -> inspect.Signature: @@ -144,9 +143,6 @@ def upgrade_data_config(self) -> dict[str, Any]: self.old_config["dataset"], ) - # Input size is a list in the old config, convert it to a tuple - init_args["image_size"] = to_tuple(init_args["image_size"]) - return { "data": { "class_path": class_path,