Refactor and restructure anomalib.data

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,33 @@
 
 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 .predict import PredictDataset
 from .utils import LabelName
-from .video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat
 
 logger = logging.getLogger(__name__)
 
@@ -63,7 +84,24 @@ 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",
+    "PredictDataset",
+    # Anomalib datamodules
     "AnomalibDataModule",
     "DepthDataFormat",
     "ImageDataFormat",
@@ -72,7 +110,6 @@ def get_datamodule(config: DictConfig | ListConfig | dict) -> AnomalibDataModule
     "BTech",
     "Folder",
     "Folder3D",
-    "PredictDataset",
     "Kolektor",
     "MVTec",
     "MVTec3D",