Refactor and restructure anomalib.data (#2302)

* Move datamodules to datamodule sub-package

* Move datamodules to datamodule sub-package

* Split datamodules and datasets

* Restructure dataclasses to data

* Fix relative imports

* Use absolute imports

* Add datasets dir

* Add relative imports for torch datasets

* Update src/anomalib/data/datamodules/base/__init__.py

Co-authored-by: Ashwin Vaidya <[email protected]>

---------

Co-authored-by: Ashwin Vaidya <[email protected]>

Author: 2 people

.gitignore

@@ -1,9 +1,8 @@
 # Project related
 datasets
+!src/anomalib/data/datasets
 pre_trained
-!anomalib/datasets
 results
-!anomalib/core/results
 
 # Test-related files and directories
 tmp*

notebooks/100_datamodules/101_btech.ipynb

@@ -71,7 +71,7 @@
     "from torchvision.transforms.v2 import Resize\n",
     "from torchvision.transforms.v2.functional import to_pil_image\n",
     "\n",
-    "from anomalib.data.image.btech import BTech, BTechDataset\n",
+    "from anomalib.data import BTech, BTechDataset\n",
     "from anomalib import TaskType"
    ]
   },

notebooks/100_datamodules/102_mvtec.ipynb

@@ -33,7 +33,7 @@
     "from torchvision.transforms.v2 import Resize\n",
     "from torchvision.transforms.v2.functional import to_pil_image\n",
     "\n",
-    "from anomalib.data.image.mvtec import MVTec, MVTecDataset\n",
+    "from anomalib.data import MVTec, MVTecDataset\n",
     "from anomalib import TaskType"
    ]
   },

notebooks/100_datamodules/103_folder.ipynb

@@ -73,7 +73,7 @@
     "from torchvision.transforms.v2 import Resize\n",
     "from torchvision.transforms.v2.functional import to_pil_image\n",
     "\n",
-    "from anomalib.data.image.folder import Folder, FolderDataset\n",
+    "from anomalib.data import Folder, FolderDataset\n",
     "from anomalib import TaskType"
    ]
   },

src/anomalib/callbacks/metrics.py

@@ -13,7 +13,7 @@
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 
 from anomalib import TaskType
-from anomalib.dataclasses import Batch
+from anomalib.data import Batch
 from anomalib.metrics import AnomalibMetricCollection, create_metric_collection
 from anomalib.models import AnomalyModule
 

src/anomalib/data/__init__.py

@@ -12,12 +12,36 @@
 
 from anomalib.utils.config import to_tuple
 
-from .base import AnomalibDataModule, AnomalibDataset
-from .depth import DepthDataFormat, Folder3D, MVTec3D
-from .image import BTech, Folder, ImageDataFormat, Kolektor, MVTec, Visa
+# Dataclasses
+from .dataclasses import (
+    Batch,
+    DatasetItem,
+    DepthBatch,
+    DepthItem,
+    ImageBatch,
+    ImageItem,
+    InferenceBatch,
+    NumpyImageBatch,
+    NumpyImageItem,
+    NumpyVideoBatch,
+    NumpyVideoItem,
+    VideoBatch,
+    VideoItem,
+)
+
+# Datamodules
+from .datamodules.base import AnomalibDataModule
+from .datamodules.depth import DepthDataFormat, Folder3D, MVTec3D
+from .datamodules.image import BTech, Folder, ImageDataFormat, Kolektor, MVTec, Visa
+from .datamodules.video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat
+
+# Datasets
+from .datasets import AnomalibDataset
+from .datasets.depth import Folder3DDataset, MVTec3DDataset
+from .datasets.image import BTechDataset, FolderDataset, KolektorDataset, MVTecDataset, VisaDataset
+from .datasets.video import AvenueDataset, ShanghaiTechDataset, UCSDpedDataset
 from .predict import PredictDataset
 from .utils import LabelName
-from .video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat
 
 logger = logging.getLogger(__name__)
 
@@ -63,7 +87,34 @@ def get_datamodule(config: DictConfig | ListConfig | dict) -> AnomalibDataModule
 
 
 __all__ = [
+    # Anomalib dataclasses
+    "DatasetItem",
+    "Batch",
+    "InferenceBatch",
+    "ImageItem",
+    "ImageBatch",
+    "VideoItem",
+    "VideoBatch",
+    "DepthItem",
+    "DepthBatch",
+    "NumpyImageItem",
+    "NumpyImageBatch",
+    "NumpyVideoItem",
+    "NumpyVideoBatch",
+    # Anomalib datasets
     "AnomalibDataset",
+    "Folder3DDataset",
+    "MVTec3DDataset",
+    "BTechDataset",
+    "FolderDataset",
+    "KolektorDataset",
+    "MVTecDataset",
+    "VisaDataset",
+    "AvenueDataset",
+    "ShanghaiTechDataset",
+    "UCSDpedDataset",
+    "PredictDataset",
+    # Anomalib datamodules
     "AnomalibDataModule",
     "DepthDataFormat",
     "ImageDataFormat",
@@ -72,7 +123,6 @@ def get_datamodule(config: DictConfig | ListConfig | dict) -> AnomalibDataModule
     "BTech",
     "Folder",
     "Folder3D",
-    "PredictDataset",
     "Kolektor",
     "MVTec",
     "MVTec3D",

src/anomalib/data/base/__init__.py

@@ -53,17 +53,19 @@
 )
 
 __all__ = [
+    # Numpy
+    "NumpyImageItem",
+    "NumpyImageBatch",
+    "NumpyVideoItem",
+    "NumpyVideoBatch",
+    # Torch
     "DatasetItem",
     "Batch",
     "InferenceBatch",
     "ImageItem",
     "ImageBatch",
     "VideoItem",
     "VideoBatch",
-    "NumpyImageItem",
-    "NumpyImageBatch",
-    "NumpyVideoItem",
-    "NumpyVideoBatch",
     "DepthItem",
     "DepthBatch",
 ]

src/anomalib/dataclasses/__init__.py src/anomalib/data/dataclasses/__init__.pysrc/anomalib/dataclasses/__init__.py renamed to src/anomalib/data/dataclasses/__init__.py

@@ -0,0 +1,24 @@
+"""Numpy-based dataclasses for Anomalib.
+
+This module provides numpy-based implementations of the generic dataclasses
+used in Anomalib. These classes are designed to work with numpy arrays for
+efficient data handling and processing in anomaly detection tasks.
+
+The module includes the following main classes:
+
+- NumpyItem: Represents a single item in Anomalib datasets using numpy arrays.
+- NumpyBatch: Represents a batch of items in Anomalib datasets using numpy arrays.
+- NumpyImageItem: Represents a single image item with additional image-specific fields.
+- NumpyImageBatch: Represents a batch of image items with batch operations.
+- NumpyVideoItem: Represents a single video item with video-specific fields.
+- NumpyVideoBatch: Represents a batch of video items with video-specific operations.
+"""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .base import NumpyBatch, NumpyItem
+from .image import NumpyImageBatch, NumpyImageItem
+from .video import NumpyVideoBatch, NumpyVideoItem
+
+__all__ = ["NumpyBatch", "NumpyItem", "NumpyImageBatch", "NumpyImageItem", "NumpyVideoBatch", "NumpyVideoItem"]