diff --git a/docs/source/markdown/guides/how_to/data/custom_data.md b/docs/source/markdown/guides/how_to/data/custom_data.md index 7370db205a..8f3161fb9e 100644 --- a/docs/source/markdown/guides/how_to/data/custom_data.md +++ b/docs/source/markdown/guides/how_to/data/custom_data.md @@ -1,4 +1,4 @@ -# Custom Data +# Custom Data (Images) This tutorial will show you how to train anomalib models on your custom data. More specifically, we will show you how to use the [Folder](../../reference/data/image/folder.md) diff --git a/docs/source/markdown/guides/how_to/data/custom_data_videos.md b/docs/source/markdown/guides/how_to/data/custom_data_videos.md new file mode 100644 index 0000000000..45529707c3 --- /dev/null +++ b/docs/source/markdown/guides/how_to/data/custom_data_videos.md @@ -0,0 +1,347 @@ +# Custom Data (Videos) + +This tutorial will show you how to train anomalib models on your custom +data. More specifically, we will show you how to use the [FolderVideo](../../reference/data/video/folder_video.md) +dataset to train anomalib models on your custom data. + +```{warning} +This tutorial assumes that you have already installed anomalib. +If not, please refer to the installation section. +``` + +```{note} +We will use our [MicroDatasets](https://github.com/openvinotoolkit/anomalib/releases/download/microvideodatasets/microdatasets.zip) three datasets to show how the folder structure should be for the compatibility of the [FolderVideo](../../reference/data/video/folder_video.md) +dataset, but you can use any dataset you want. +``` + +```{note} +The folder structure should be: +- path/to/dataset/ <-- root + - folder_containing_normal_videos/ <- normal_dir + - folder_containing_abnormal_videos/ <- test_dir + - folder_containing_masks_for_abnormal_videos/ <- mask_dir + +If the videos are in a frame by frame format, every video should be inside a different folder in test_dir: + +- test_dir/ + - vid_00/ + - frame_00.png + - frame_01.png + - frame_02.png + -... + - vid_01/ + - frame_00.png + - frame_01.png + - frame_02.png + -... +``` + +We will split the section to two tasks: Classification and Segmentation. + +## Classification Dataset + +In certain use-cases, ground-truth masks for the abnormal Frames of the video may not be +available. In such cases, we could use the classification task to train a model +that will be able to detect the abnormal Frames in the test set. + +We will split this section into two tasks: + +- Classification with normal and abnormal frame masks. +- Classification with only normal videos. + +### With Normal and Abnormal frame masks + +We could use [FolderVideo](../../reference/data/video/folder_video.md) datamodule to train +a model on this dataset. We could run the following python code to create the +custom datamodule: + +:::::{dropdown} Code Syntax +:icon: code + +::::{tab-set} +:::{tab-item} API +:sync: label-1 + +```{literalinclude} ../../../../snippets/data/video/folder_video/classification/default.txt +:language: python +``` + +```{note} +As can be seen above, we only need to specify the ``task`` argument to ``classification``. We could have also use ``TaskType.CLASSIFICATION`` instead of ``classification``. +``` + +The [FolderVideo](../../reference/data/video/folder_video.md) datamodule will create training, validation, test and +prediction datasets and dataloaders for us. We can access the datasets +and dataloaders by following the same approach as in the segmentation +task. + +When we check the samples from the dataloaders, we will see that the +`mask` key is not present in the samples. This is because we do not need +the masks for the classification task. + +```{literalinclude} ../../../../snippets/data/video/folder_video/classification/dataloader_values.txt +:language: python +``` + +Training the model is as simple as running the following command: + +```{literalinclude} ../../../../snippets/train/api/classification/model_and_engine_video.txt +:language: python +``` + +where we train a AiVad model on this custom dataset with default model parameters. + +::: + +:::{tab-item} CLI +:sync: label-2 + +Here is the dataset config to create the same custom datamodule: + +```{literalinclude} ../../../../snippets/config/data/video/folder/classification/cli/default.yaml +:language: yaml +``` + +Assume that we have saved the above config file as `classification.yaml`. We could run the following CLI command to train a AiVad model on above dataset: + +```bash +anomalib train --data classification.yaml --model anomalib.models.AiVad --task CLASSIFICATION +``` + +```{note} +As can be seen above, we also need to specify the ``task`` argument to ``CLASSIFICATION`` to explicitly tell anomalib that we want to train a classification model. This is because the default `task` is `SEGMENTATION` within `Engine`. +``` + +::: +:::: +::::: + +### With Only Normal Videos + +There are certain cases where we only have normal frames in our dataset but +would like to train a classification model. + +This could be done in two ways: + +- Train the model and skip the validation and test steps, as we do not have + abnormal Videos to validate and test the model on. +- Use the synthetic anomaly generation feature to create abnormal videos/masks from + normal frames, and perform the validation and test steps. + +For now we will focus on the second approach. + +#### With Validation and Testing via Synthetic Anomalies + +If we want to check the performance of the model, we will need to have abnormal +videos/masks to validate and test the model on. During the validation stage, +these anomalous videos are used to normalize the anomaly scores and find the +best threshold that separates normal and abnormal videos. + +Anomalib provides synthetic anomaly generation capabilities to create abnormal +frames from normal frames so we could check the performance. We could use the +[FolderVideo](../../reference/data/video/folder_video.md) datamodule to train a model on +this dataset. + +:::::{dropdown} Code Syntax +:icon: code + +::::{tab-set} +:::{tab-item} API +:sync: label-1 + +```{literalinclude} ../../../../snippets/data/video/folder_video/classification/normal_and_synthetic.txt +:language: python +``` + +Once the datamodule is setup, the rest of the process is the same as in the previous classification example. + +```{literalinclude} ../../../../snippets/train/api/classification/model_and_engine_video.txt +:language: python +``` + +where we train a AiVad model on this custom dataset with default model parameters. + +::: + +:::{tab-item} CLI +:sync: label-2 + +Here is the CLI command to create the same custom datamodule with only normal videos. We only need to change the `test_split_mode` argument to `SYNTHETIC` to generate synthetic anomalies. + +```{literalinclude} ../../../../snippets/config/data/video/folder/classification/cli/normal_and_synthetic.yaml +:language: yaml +``` + +Assume that we have saved the above config file as `normal.yaml`. We could run the following CLI command to train a AiVad model on above dataset: + +```bash +anomalib train --data normal.yaml --model anomalib.models.AiVad --task CLASSIFICATION +``` + +```{note} +As shown in the previous classification example, we, again, need to specify the ``task`` argument to ``CLASSIFICATION`` to explicitly tell anomalib that we want to train a classification model. This is because the default `task` is `SEGMENTATION` within `Engine`. +``` + +::: +:::: +::::: + +## Segmentation Dataset + +Assume that we have a dataset in which the training set contains only +normal videos, and the test set contains both normal and abnormal +videos. We also have masks for the abnormal video-frames in the test set. We +want to train an anomaly segmentation model that will be able to detect the +abnormal regions in the test set. + +### With Normal and Abnormal Videos + +We could use [FolderVideo](../../reference/data/video/folder_video.md) datamodule to load the microvideo dataset in a format that is readable by Anomalib's models. + +:::::{dropdown} Code Syntax +::::{tab-set} +:::{tab-item} API +:sync: label-1 + +We could run the following python code to create the custom datamodule: + +```{literalinclude} ../../../../snippets/data/video/folder_video/segmentation/default.txt +:language: python +``` + +The [FolderVideo](../../reference/data/video/folder_video.md) datamodule will create training, validation, test and +prediction datasets and dataloaders for us. We can access the datasets +and dataloaders using the following attributes: + +```{literalinclude} ../../../../snippets/data/video/folder_video/segmentation/datamodule_attributes.txt +:language: python +``` + +To check what individual samples from dataloaders look like, we can run +the following command: + +```{literalinclude} ../../../../snippets/data/video/folder_video/segmentation/dataloader_values.txt +:language: python +``` + +We could check the shape of the videos and masks using the following +commands: + +```python +print(train_data["image"].shape) +# torch.Size([2, 2, 3, 360, 640]) + +print(train_data["mask"].shape) +# torch.Size([2, 360, 640]) +``` + +Training the model is as simple as running the following command: + +```{literalinclude} ../../../../snippets/train/api/segmentation/model_and_engine_video.txt +:language: python +``` + +where we train a AiVad model on this custom dataset with default model parameters. +::: + +:::{tab-item} CLI +:sync: label-2 +Here is the CLI command to create the same custom datamodule: + +```{literalinclude} ../../../../snippets/config/data/video/folder/segmentation/cli/default.yaml +:language: yaml +``` + +Assume that we have saved the above config file as `segmentation.yaml`. +We could run the following CLI command to train a AiVad model on above dataset: + +```bash +anomalib train --data segmentation.yaml --model anomalib.models.AiVad +``` + +::: + +:::: +::::: + +This example demonstrates how to create a segmentation dataset with +normal and abnormal videos. We could expand this example to create a +segmentation dataset with only normal videos. + +### With Only Normal Videos + +There are certain cases where we only have normal videos in our dataset +but would like to train a segmentation model. This could be done in two ways: + +- Train the model and skip the validation and test steps, as + we do not have abnormal videos to validate and test the model on, or +- Use the synthetic anomaly generation feature to create abnormal + videos from normal videos, and perform the validation and test steps. + +For now we will focus on the second approach. + +#### With Validation and Testing via Synthetic Anomalies + +We could use the synthetic anomaly generation feature again to create abnormal +videos/frames from normal videos. We could then use the +[FolderVideo](../../reference/data/video/folder_video.md) datamodule to train a model on +this dataset. Here is the python code to create the custom datamodule: + +:::::{dropdown} Code Syntax +:icon: code + +::::{tab-set} +:::{tab-item} API +:sync: label-1 + +We could run the following python code to create the custom datamodule: + +```{literalinclude} ../../../../snippets/data/video/folder_video/segmentation/normal_and_synthetic.txt +:language: python +``` + +As can be seen from the code above, we only need to specify the +`test_split_mode` argument to `SYNTHETIC`. The [FolderVideo](../../reference/data/video/folder_video.md) datamodule will create training, validation, test and prediction datasets and +dataloaders for us. + +To check what individual samples from dataloaders look like, we can run +the following command: + +```{literalinclude} ../../../../snippets/data/video/folder_video/segmentation/dataloader_values.txt +:language: python +``` + +We could check the shape of the videos and masks using the following +commands: + +```python +print(train_data["image"].shape) +# torch.Size([2, 2, 3, 360, 640]) + +print(train_data["mask"].shape) +# torch.Size([2, 360, 640]) +``` + +::: + +:::{tab-item} CLI +:sync: label-2 + +Here is the CLI command to create the same custom datamodule with only normal +videos. We only need to change the `test_split_mode` argument to `SYNTHETIC` to +generate synthetic anomalies. + +```{literalinclude} ../../../../snippets/config/data/video/folder/segmentation/cli/normal_and_synthetic.yaml +:language: yaml +``` + +Assume that we have saved the above config file as `synthetic.yaml`. We could +run the following CLI command to train a AiVad model on above dataset: + +```bash +anomalib train --data synthetic.yaml --model anomalib.models.AiVad +``` + +::: +:::: +::::: diff --git a/docs/source/markdown/guides/how_to/data/index.md b/docs/source/markdown/guides/how_to/data/index.md index 19b95a9eae..407e012bac 100644 --- a/docs/source/markdown/guides/how_to/data/index.md +++ b/docs/source/markdown/guides/how_to/data/index.md @@ -6,11 +6,18 @@ This section contains tutorials on how to fully utilize the data components of a :margin: 1 1 0 0 :gutter: 1 -:::{grid-item-card} {octicon}`database` Train on Custom Data. +:::{grid-item-card} {octicon}`database` Train on Custom Data(Images). :link: ./custom_data :link-type: doc -Learn more about how to use `Folder` dataset to train anomalib models on your custom data. +Learn more about how to use `Folder` dataset to train anomalib models on your custom data (Images). +::: + +:::{grid-item-card} {octicon}`database` Train on Custom Data(Videos). +:link: ./custom_data_videos +:link-type: doc + +Learn more about how to use `FolderVideo` dataset to train anomalib models on your custom data (Videos). ::: :::{grid-item-card} {octicon}`table` Input tiling @@ -27,5 +34,6 @@ Learn more about how to use the tiler for input tiling. :hidden: ./custom_data +./custom_data_videos ./input_tiling ``` diff --git a/docs/source/markdown/guides/reference/data/video/folder_video.md b/docs/source/markdown/guides/reference/data/video/folder_video.md new file mode 100644 index 0000000000..6e67788df6 --- /dev/null +++ b/docs/source/markdown/guides/reference/data/video/folder_video.md @@ -0,0 +1,7 @@ +# FolderVideo Data + +```{eval-rst} +.. automodule:: anomalib.data.video.folder_video + :members: + :show-inheritance: +``` diff --git a/docs/source/markdown/guides/reference/data/video/index.md b/docs/source/markdown/guides/reference/data/video/index.md index 9f357053fa..3e828496b5 100644 --- a/docs/source/markdown/guides/reference/data/video/index.md +++ b/docs/source/markdown/guides/reference/data/video/index.md @@ -23,6 +23,13 @@ Learn more about Shanghai Tech dataset. Learn more about UCSD Ped1 and Ped2 datasets. ::: +:::{grid-item-card} FolderVideo +:link: ./folder_video +:link-type: doc + +Learn more about custom folder dataset(video). +::: + :::: ```{toctree} @@ -30,6 +37,7 @@ Learn more about UCSD Ped1 and Ped2 datasets. :hidden: ./avenue +./folder_video ./shanghaitech ./ucsd_ped ``` diff --git a/docs/source/snippets/config/data/video/folder/classification/cli/default.yaml b/docs/source/snippets/config/data/video/folder/classification/cli/default.yaml new file mode 100644 index 0000000000..0f9a706311 --- /dev/null +++ b/docs/source/snippets/config/data/video/folder/classification/cli/default.yaml @@ -0,0 +1,16 @@ +class_path: anomalib.data.FolderVideo +init_args: + root: "datasets/microvideodataset/CustomDataset_1" + normal_dir: "train_vid" + mask_dir: "test_labels" + test_dir: "test_vid" + image_size: [256, 256] + train_batch_size: 2 + eval_batch_size: 2 + num_workers: 0 + task: CLASSIFICATION + train_transform: null + eval_transform: null + val_split_mode: SAME_AS_TEST + val_split_ratio: 0.5 + seed: null diff --git a/docs/source/snippets/config/data/video/folder/classification/cli/normal.yaml b/docs/source/snippets/config/data/video/folder/classification/cli/normal.yaml new file mode 100644 index 0000000000..4bcaa3cce4 --- /dev/null +++ b/docs/source/snippets/config/data/video/folder/classification/cli/normal.yaml @@ -0,0 +1,8 @@ +class_path: anomalib.data.Folder +init_args: + root: "datasets/hazelnut_toy" + normal_dir: "good" + normalization: imagenet + val_split_mode: NONE + test_split_mode: NONE + task: classification diff --git a/docs/source/snippets/config/data/video/folder/classification/cli/normal_and_synthetic.yaml b/docs/source/snippets/config/data/video/folder/classification/cli/normal_and_synthetic.yaml new file mode 100644 index 0000000000..f17c652739 --- /dev/null +++ b/docs/source/snippets/config/data/video/folder/classification/cli/normal_and_synthetic.yaml @@ -0,0 +1,9 @@ +class_path: anomalib.data.FolderVideo +init_args: + root: "datasets/microvideodataset/CustomDataset_1" + normal_dir: "train_vid" + test_split_mode: SYNTHETIC + task: CLASSIFICATION + train_batch_size: 2 + eval_batch_size: 2 + num_workers: 0 diff --git a/docs/source/snippets/config/data/video/folder/segmentation/cli/default.yaml b/docs/source/snippets/config/data/video/folder/segmentation/cli/default.yaml new file mode 100644 index 0000000000..9998b34ed3 --- /dev/null +++ b/docs/source/snippets/config/data/video/folder/segmentation/cli/default.yaml @@ -0,0 +1,7 @@ +class_path: anomalib.data.FolderVideo +init_args: + root: "datasets/microvideodataset/CustomDataset_1" + normal_dir: "train_vid" + mask_dir: "test_labels" + test_dir: "test_vid" + normal_split_ratio: 0.2 diff --git a/docs/source/snippets/config/data/video/folder/segmentation/cli/normal.yaml b/docs/source/snippets/config/data/video/folder/segmentation/cli/normal.yaml new file mode 100644 index 0000000000..5fe8211a29 --- /dev/null +++ b/docs/source/snippets/config/data/video/folder/segmentation/cli/normal.yaml @@ -0,0 +1,10 @@ +class_path: anomalib.data.Folder +init_args: + root: "datasets/hazelnut_toy" + normal_dir: "good" + abnormal_dir: "crack" + mask_dir: "mask/crack" + normalization: imagenet + test_split_mode: NONE + val_split_mode: NONE + seed: null diff --git a/docs/source/snippets/config/data/video/folder/segmentation/cli/normal_and_synthetic.yaml b/docs/source/snippets/config/data/video/folder/segmentation/cli/normal_and_synthetic.yaml new file mode 100644 index 0000000000..71a16ef8ba --- /dev/null +++ b/docs/source/snippets/config/data/video/folder/segmentation/cli/normal_and_synthetic.yaml @@ -0,0 +1,15 @@ +class_path: anomalib.data.FolderVideo +init_args: + root: "datasets/microvideodataset/CustomDataset_1" + normal_dir: "train_vid" + train_batch_size: 2 + eval_batch_size: 2 + num_workers: 0 + task: SEGMENTATION + train_transform: null + eval_transform: null + test_split_mode: SYNTHETIC + test_split_ratio: 0.2 + val_split_mode: same_as_test + val_split_ratio: 0.5 + seed: null diff --git a/docs/source/snippets/data/video/folder_video/classification/dataloader_values.txt b/docs/source/snippets/data/video/folder_video/classification/dataloader_values.txt new file mode 100644 index 0000000000..11e8e2f6ef --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/classification/dataloader_values.txt @@ -0,0 +1,11 @@ +i, train_data = next(enumerate(datamodule.train_dataloader())) +print(train_data.keys()) +# dict_keys(['image', 'video_path', 'frames', 'last_frame', 'original_image']) + +i, val_data = next(enumerate(datamodule.val_dataloader())) +print(val_data.keys()) +# dict_keys(['image', 'video_path', 'frames', 'last_frame', 'original_image']) + +i, test_data = next(enumerate(datamodule.test_dataloader())) +print(test_data.keys()) +# dict_keys(['image', 'video_path', 'frames', 'last_frame', 'original_image']) diff --git a/docs/source/snippets/data/video/folder_video/classification/datamodule_attributes.txt b/docs/source/snippets/data/video/folder_video/classification/datamodule_attributes.txt new file mode 100644 index 0000000000..5e5408e27f --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/classification/datamodule_attributes.txt @@ -0,0 +1,9 @@ +# Access the datasets +train_dataset = datamodule.train_data +val_dataset = datamodule.val_data +test_dataset = datamodule.test_data + +# Access the dataloaders +train_dataloader = datamodule.train_dataloader() +val_dataloader = datamodule.val_dataloader() +test_dataloader = datamodule.test_dataloader() diff --git a/docs/source/snippets/data/video/folder_video/classification/default.txt b/docs/source/snippets/data/video/folder_video/classification/default.txt new file mode 100644 index 0000000000..f546137419 --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/classification/default.txt @@ -0,0 +1,16 @@ +# Import the datamodule +from anomalib.data import FolderVideo + +# Create the datamodule with a reduced batch size +datamodule = FolderVideo( + root="datasets/microvideodataset/CustomDataset_1", + normal_dir="train_vid", + test_dir="test_vid", + task="classification", + train_batch_size=2, + eval_batch_size=2, + num_workers=0, +) + +# Setup the datamodule +datamodule.setup() diff --git a/docs/source/snippets/data/video/folder_video/classification/normal.txt b/docs/source/snippets/data/video/folder_video/classification/normal.txt new file mode 100644 index 0000000000..01355d5b8e --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/classification/normal.txt @@ -0,0 +1,16 @@ +# Import the datamodule +from anomalib.data import Folder +from anomalib.data.utils import TestSplitMode, ValSplitMode + +# Create the datamodule +datamodule = Folder( + name="hazelnut_toy", + root="datasets/hazelnut_toy", + normal_dir="good", + val_split_mode=ValSplitMode.NONE, + test_split_mode=TestSplitMode.NONE, + task="classification", +) + +# Setup the datamodule +datamodule.setup() diff --git a/docs/source/snippets/data/video/folder_video/classification/normal_and_synthetic.txt b/docs/source/snippets/data/video/folder_video/classification/normal_and_synthetic.txt new file mode 100644 index 0000000000..af19c4a944 --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/classification/normal_and_synthetic.txt @@ -0,0 +1,17 @@ +# Import the datamodule +from anomalib.data import FolderVideo +from anomalib.data.utils import TestSplitMode + +# Create the datamodule +datamodule = FolderVideo( + root="datasets/microvideodataset/CustomDataset_1", + test_split_mode=TestSplitMode.SYNTHETIC, + normal_dir="train_vid", + task="classification", + train_batch_size=2, + eval_batch_size=2, + num_workers=0, +) + +# Setup the datamodule +datamodule.setup() diff --git a/docs/source/snippets/data/video/folder_video/segmentation/dataloader_values.txt b/docs/source/snippets/data/video/folder_video/segmentation/dataloader_values.txt new file mode 100644 index 0000000000..185fe19871 --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/segmentation/dataloader_values.txt @@ -0,0 +1,11 @@ +i, train_data = next(enumerate(datamodule.train_dataloader())) +print(train_data.keys()) +dict_keys(['image', 'video_path', 'frames', 'last_frame', 'original_image']) + +i, val_data = next(enumerate(datamodule.val_dataloader())) +print(val_data.keys()) +dict_keys(['image', 'mask', 'video_path', 'frames', 'last_frame', 'original_image', 'label']) + +i, test_data = next(enumerate(datamodule.test_dataloader())) +print(test_data.keys()) +dict_keys(['image', 'mask', 'video_path', 'frames', 'last_frame', 'original_image', 'label']) diff --git a/docs/source/snippets/data/video/folder_video/segmentation/datamodule_attributes.txt b/docs/source/snippets/data/video/folder_video/segmentation/datamodule_attributes.txt new file mode 100644 index 0000000000..5e5408e27f --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/segmentation/datamodule_attributes.txt @@ -0,0 +1,9 @@ +# Access the datasets +train_dataset = datamodule.train_data +val_dataset = datamodule.val_data +test_dataset = datamodule.test_data + +# Access the dataloaders +train_dataloader = datamodule.train_dataloader() +val_dataloader = datamodule.val_dataloader() +test_dataloader = datamodule.test_dataloader() diff --git a/docs/source/snippets/data/video/folder_video/segmentation/default.txt b/docs/source/snippets/data/video/folder_video/segmentation/default.txt new file mode 100644 index 0000000000..47d9785370 --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/segmentation/default.txt @@ -0,0 +1,16 @@ +# Import the datamodule +from anomalib.data import FolderVideo + +# Create the datamodule with a reduced batch size +datamodule = FolderVideo( + root="datasets/microvideodataset/CustomDataset_1", + normal_dir="train_vid", + test_dir="test_vid", + mask_dir="test_labels", + train_batch_size=2, + eval_batch_size=2, + num_workers=0, +) + +# Setup the datamodule +datamodule.setup() diff --git a/docs/source/snippets/data/video/folder_video/segmentation/normal.txt b/docs/source/snippets/data/video/folder_video/segmentation/normal.txt new file mode 100644 index 0000000000..2d01d22674 --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/segmentation/normal.txt @@ -0,0 +1,15 @@ +# Import the datamodule +from anomalib.data import Folder +from anomalib.data.utils import TestSplitMode, ValSplitMode + +# Create the datamodule +datamodule = Folder( + name="hazelnut_toy", + root="datasets/MVTec/transistor", + normal_dir="train/good", + val_split_mode=ValSplitMode.NONE, + test_split_mode=TestSplitMode.NONE, +) + +# Setup the datamodule +datamodule.setup() diff --git a/docs/source/snippets/data/video/folder_video/segmentation/normal_and_synthetic.txt b/docs/source/snippets/data/video/folder_video/segmentation/normal_and_synthetic.txt new file mode 100644 index 0000000000..a238447b0c --- /dev/null +++ b/docs/source/snippets/data/video/folder_video/segmentation/normal_and_synthetic.txt @@ -0,0 +1,14 @@ +# Import the datamodule +from anomalib.data import Folder +from anomalib.data.utils import TestSplitMode + +# Create the datamodule +datamodule = Folder( + name="hazelnut_toy", + root="datasets/hazelnut_toy", + normal_dir="good", + test_split_mode=TestSplitMode.SYNTHETIC, +) + +# Setup the datamodule +datamodule.setup() diff --git a/docs/source/snippets/train/api/classification/model_and_engine_video.txt b/docs/source/snippets/train/api/classification/model_and_engine_video.txt new file mode 100644 index 0000000000..25a8c3e0a7 --- /dev/null +++ b/docs/source/snippets/train/api/classification/model_and_engine_video.txt @@ -0,0 +1,10 @@ +# Import the model and engine +from anomalib.models import AiVad +from anomalib.engine import Engine + +# Create the model and engine +model = AiVad() +engine = Engine(task="classification") + +# Train a Patchcore model on the given datamodule +engine.train(datamodule=datamodule, model=model) diff --git a/docs/source/snippets/train/api/segmentation/model_and_engine_video.txt b/docs/source/snippets/train/api/segmentation/model_and_engine_video.txt new file mode 100644 index 0000000000..f17dabe889 --- /dev/null +++ b/docs/source/snippets/train/api/segmentation/model_and_engine_video.txt @@ -0,0 +1,10 @@ +# Import the model and engine +from anomalib.models import AiVad +from anomalib.engine import Engine + +# Create the model and engine +model = AiVad() +engine = Engine() + +# Train a Patchcore model on the given datamodule +engine.train(datamodule=datamodule, model=model) diff --git a/src/anomalib/data/__init__.py b/src/anomalib/data/__init__.py index 85a4fd1589..85ea614bae 100644 --- a/src/anomalib/data/__init__.py +++ b/src/anomalib/data/__init__.py @@ -3,7 +3,6 @@ # Copyright (C) 2022-2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 - import importlib import logging from enum import Enum @@ -18,7 +17,7 @@ from .image import BTech, Folder, ImageDataFormat, Kolektor, MVTec, Visa from .predict import PredictDataset from .utils import LabelName -from .video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat +from .video import Avenue, FolderVideo, ShanghaiTech, UCSDped, VideoDataFormat logger = logging.getLogger(__name__) @@ -64,6 +63,7 @@ def get_datamodule(config: DictConfig | ListConfig) -> AnomalibDataModule: "MVTec", "MVTec3D", "Avenue", + "FolderVideo", "UCSDped", "ShanghaiTech", "Visa", diff --git a/src/anomalib/data/utils/video.py b/src/anomalib/data/utils/video.py index 7a939ea861..1992c5ec49 100644 --- a/src/anomalib/data/utils/video.py +++ b/src/anomalib/data/utils/video.py @@ -5,6 +5,7 @@ import warnings from abc import ABC, abstractmethod +from collections import Counter from pathlib import Path from typing import Any @@ -98,3 +99,39 @@ def convert_video(input_path: Path, output_path: Path, codec: str = "MP4V") -> N video_reader.release() video_writer.release() + + +def most_common_extension(folder_path: Path) -> str | None: + """Determine the most common file extension in a specified folder and its subfolders. + + Args: + folder_path (Path): The path to the folder to be analyzed. + + Returns: + str: The most common file extension in the folder and its subfolders. + None: If no files with extensions are found in the folder and its subfolders. + """ + if not folder_path.is_dir(): + return folder_path.suffix + + extensions = [] + + # Check files in the main folder + extensions = [f.suffix for f in folder_path.iterdir() if f.is_file() and f.suffix] + + # Check files in immediate subfolders + ext_subfolders = [ + f.suffix + for subdir in folder_path.iterdir() + if subdir.is_dir() + for f in subdir.iterdir() + if f.is_file() and f.suffix + ] + extensions += ext_subfolders + + # Count extensions and find the most common + extension_counts = Counter(extensions) + if extension_counts: + return extension_counts.most_common(1)[0][0] + + return None diff --git a/src/anomalib/data/video/__init__.py b/src/anomalib/data/video/__init__.py index a9651529bf..c58ab3b06b 100644 --- a/src/anomalib/data/video/__init__.py +++ b/src/anomalib/data/video/__init__.py @@ -3,10 +3,10 @@ # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 - from enum import Enum from .avenue import Avenue +from .folder_video import FolderVideo from .shanghaitech import ShanghaiTech from .ucsd_ped import UCSDped @@ -17,6 +17,7 @@ class VideoDataFormat(str, Enum): UCSDPED = "ucsdped" AVENUE = "avenue" SHANGHAITECH = "shanghaitech" + FOLDERVIDEO = "foldervideo" -__all__ = ["Avenue", "ShanghaiTech", "UCSDped"] +__all__ = ["Avenue", "ShanghaiTech", "UCSDped", "FolderVideo"] diff --git a/src/anomalib/data/video/avenue.py b/src/anomalib/data/video/avenue.py index 831caa4021..11fd7473ef 100644 --- a/src/anomalib/data/video/avenue.py +++ b/src/anomalib/data/video/avenue.py @@ -49,12 +49,12 @@ DATASET_DOWNLOAD_INFO = DownloadInfo( name="Avenue Dataset", - url="http://www.cse.cuhk.edu.hk/leojia/projects/detectabnormal/Avenue_Dataset.zip", + url="https://www.cse.cuhk.edu.hk/leojia/projects/detectabnormal/Avenue_Dataset.zip", hashsum="fc9cb8432a11ca79c18aa180c72524011411b69d3b0ff27c8816e41c0de61531", ) ANNOTATIONS_DOWNLOAD_INFO = DownloadInfo( name="Avenue Annotations", - url="http://www.cse.cuhk.edu.hk/leojia/projects/detectabnormal/ground_truth_demo.zip", + url="https://www.cse.cuhk.edu.hk/leojia/projects/detectabnormal/ground_truth_demo.zip", hashsum="60fec1728ec8f73a58aad3aeb5729d70a805a47e0b8eb4bf91ab67ef06386d77", ) diff --git a/src/anomalib/data/video/folder_video.py b/src/anomalib/data/video/folder_video.py new file mode 100644 index 0000000000..a0774adfb4 --- /dev/null +++ b/src/anomalib/data/video/folder_video.py @@ -0,0 +1,423 @@ +"""Custom video Folder Dataset. + +This script creates a custom dataset from a folder. +""" + +# Copyright (C) 2022-2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from pandas import DataFrame +from torchvision.transforms.v2 import Transform + +from anomalib import TaskType +from anomalib.data.base import AnomalibVideoDataModule, AnomalibVideoDataset +from anomalib.data.base.video import VideoTargetFrame +from anomalib.data.utils import ( + Split, + ValSplitMode, + read_image, + read_mask, + validate_path, +) +from anomalib.data.utils.video import ClipsIndexer, convert_video, most_common_extension + +logger = logging.getLogger(__name__) + +SINGLE_FILE_EXTENSIONS: set[str] = {".avi", ".mp4", ".npy", ".pt"} +FOLDER_IMAGE_EXTENSIONS: set[str] = {".bmp", ".png", ".jpg", ".tiff", ".tif"} + + +def make_folder_video_dataset( + root: str | Path, + normal_dir: str | Path = "", + mask_dir: str | Path | None = "", + test_dir: str | Path = "", + split: str | Split | None = None, +) -> DataFrame: + """Make Folder Video Dataset. + + Args: + root (str | Path ): Path to the root directory of the dataset. + normal_dir (str | Path ): Path to the directory containing normal images. + test_dir (str | Path, optional): Path to the directory containing abnormal images. + Defaults to ``""``. + mask_dir (str | Path, optional): Path to the directory containing the mask annotations. + Defaults to ``""``. + split (str | Split | None, optional): Dataset split (ie., Split.FULL, Split.TRAIN or Split.TEST). + Defaults to ``None``. + + Return: + DataFrame: an output dataframe containing samples for the requested split (ie., train or test). + + Example: + The following example shows how to get testing samples from a custom dataset: + + >>> root = Path('./myDataset') + >>> samples = make_folder_video_dataset( + root,normal_dir="train_vid", test_dir="test_vid", mask_dir="test_labels", split='test' + ) + >>> samples.head() + root image_path split mask_path + 0 myDataset myDataset/test_vid/01_0014 test myDataset/test_labels/01_0014.npy + 1 myDataset myDataset/test_vid/01_0015 test myDataset/test_labels/01_0015.npy + ... + """ + + def _contains_files(path: Path, extensions: set[str]) -> bool: + """Check if the path contains at least one file with a given extension in the directory or its one-level subdir. + + Args: + path (Path): The path to the folder. + extensions (list): A list of file extensions to check for. + + Returns: + bool: True if there is at least one file with a given extension, False otherwise. + """ + if path.is_dir(): + for item in path.iterdir(): + if (item.is_file() and item.suffix in extensions) or ( + item.is_dir() and any(file.suffix in extensions for file in item.iterdir() if file.is_file()) + ): + return True + return False + + def _extract_samples(root: Path, path: Path) -> list: + samples_list = [] + if most_common_extension(path) in SINGLE_FILE_EXTENSIONS: + common_extension = most_common_extension(path) + samples_list.extend( + [(str(root),) + filename.parts[-2:] for filename in sorted(path.glob(f"./*{common_extension}"))], + ) + elif most_common_extension(path) in FOLDER_IMAGE_EXTENSIONS: + samples_list.extend( + [ + (str(root),) + filename.parts[-2:] + for filename in sorted(path.glob("./*")) + if _contains_files( + path=path, + extensions=SINGLE_FILE_EXTENSIONS | FOLDER_IMAGE_EXTENSIONS, + ) + and not filename.name.startswith(".") + ], + ) + return samples_list + + if isinstance(root, str): + root = Path(root) + + if isinstance(normal_dir, str): + normal_dir = Path(normal_dir) + + if isinstance(test_dir, str): + test_dir = Path(test_dir) + + if isinstance(mask_dir, str): + mask_dir = Path(mask_dir) + + root = validate_path(root) + normal_dir = validate_path(root / normal_dir) + test_dir = validate_path(root / test_dir) + + samples_list_labels = [] + if mask_dir is not None: + mask_dir = validate_path(root / mask_dir) + samples_list_labels.extend( + [ + filename.parts[-1] + for filename in sorted(mask_dir.glob("./*")) + if ( + _contains_files(path=filename, extensions=FOLDER_IMAGE_EXTENSIONS) + and not filename.name.startswith(".") + ) + or filename.suffix in [".npy", ".pt"] + ], + ) + + samples_list = [] + samples_list.extend(_extract_samples(root, normal_dir)) + + samples_list.extend(_extract_samples(root, test_dir)) + + samples = DataFrame(samples_list, columns=["root", "folder", "image_path"]) + + # Remove DS_Store + samples = samples[~samples.loc[:, "image_path"].str.contains(".DS_Store")] + + samples["image_path"] = samples.root + "/" + samples.folder + "/" + samples.image_path + + samples.loc[samples.folder == normal_dir.parts[-1], "split"] = "train" + samples.loc[samples.folder == test_dir.parts[-1], "split"] = "test" + samples_list_labels = [str(item) for item in samples_list_labels if ".DS_Store" not in str(item)] + + if mask_dir is not None: + samples.loc[samples.folder == test_dir.parts[-1], "mask_path"] = samples_list_labels + samples["mask_path"] = str(mask_dir) + "/" + samples.mask_path.astype(str) + else: + samples.loc[samples.folder == test_dir.parts[-1], "mask_path"] = "" + samples.loc[samples.folder == normal_dir.parts[-1], "mask_path"] = "" + + if split: + samples = samples[samples.split == split] + samples = samples.reset_index(drop=True) + + return samples + + +class FolderClipsIndexerVideo(ClipsIndexer): + """Clips indexer for the test set Folder video dataset.""" + + def get_mask(self, idx: int) -> np.ndarray | torch.Tensor | None: + """Retrieve the masks from the file system.""" + video_idx, frames_idx = self.get_clip_location(idx) + mask_folder = self.mask_paths[video_idx] + if mask_folder == "": # no gt masks available for this clip + return None + frames = self.clips[video_idx][frames_idx] + common_extension = most_common_extension(Path(mask_folder)) + ret = None + + if common_extension in FOLDER_IMAGE_EXTENSIONS: + mask_frames = sorted([file for file in Path(mask_folder).glob("*") if not file.name.startswith(".")]) + mask_paths = [mask_frames[idx] for idx in frames.int()] + ret = torch.stack([read_mask(mask_path, as_tensor=True) for mask_path in mask_paths]) + elif common_extension in [".npy"]: + vid_masks = np.load(mask_folder) + ret = torch.tensor(np.take(vid_masks, frames, 0)) + elif common_extension in [".pt"]: + vid_masks = torch.load(mask_folder).numpy() + ret = torch.tensor(np.take(vid_masks, frames, 0)) + + return ret + + +class FolderClipsIndexerImgFrames(FolderClipsIndexerVideo): + """Clips indexer for the test set Folder video dataset with frames as a video.""" + + def _compute_frame_pts(self) -> None: + """Retrieve the number of frames in each video.""" + self.video_pts = [] + for video_path in self.video_paths: + n_frames = len([file for file in Path(video_path).glob("*") if not file.name.startswith(".")]) # tiff + self.video_pts.append(torch.Tensor(range(n_frames))) + + self.video_fps = [None] * len(self.video_paths) # fps information cannot be inferred from folder structure + + def get_clip(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], int]: + """Get a subclip from a list of videos. + + Args: + idx (int): index of the subclip. Must be between 0 and num_clips(). + + Returns: + video (torch.Tensor) + audio (torch.Tensor) + info (dict) + video_idx (int): index of the video in `video_paths` + """ + if idx >= self.num_clips(): + msg = f"Index {idx} out of range ({self.num_clips()} number of clips)" + raise IndexError(msg) + video_idx, clip_idx = self.get_clip_location(idx) + video_path = self.video_paths[video_idx] + clip_pts = self.clips[video_idx][clip_idx] + + # Get all the frames in 000.* ~> 999.* where * is an image format + frames = sorted([file for file in Path(video_path).glob("*") if not file.name.startswith(".")]) + + frame_paths = [frames[pt] for pt in clip_pts.int()] + video = torch.stack([read_image(frame_path, as_tensor=True) for frame_path in frame_paths]) + + return video, torch.empty((1, 0)), {}, video_idx + + +class FolderDataset(AnomalibVideoDataset): + """Folder Dataset class. + + Args: + task (TaskType): Task type, 'classification', 'detection' or 'segmentation' + split (Split): Split of the dataset, usually Split.TRAIN or Split.TEST + root (Path | str): Path to the dataset. + normal_dir (Path | str): Path to the training videos of the dataset (.avi / .mp4 / imgages as frames). + test_dir (Path | str): Path to the testing videos of the dataset. + mask_dir (Path | str | None): Path to the masks for the training videos of the dataset (.npy/.pt/ images) + clip_length_in_frames (int, optional): Number of video frames in each clip. + 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. + transform (Transform, optional): Transforms that should be applied to the input images. + Defaults to ``None``. + """ + + def __init__( + self, + task: TaskType, + split: Split, + root: Path | str, + mask_dir: Path | str | None = None, + normal_dir: Path | str = "train_dir", + test_dir: Path | str = "test_dir", + clip_length_in_frames: int = 2, + frames_between_clips: int = 1, + target_frame: VideoTargetFrame = VideoTargetFrame.LAST, + transform: Transform | None = None, + ) -> None: + super().__init__( + task=task, + clip_length_in_frames=clip_length_in_frames, + frames_between_clips=frames_between_clips, + target_frame=target_frame, + transform=transform, + ) + + self.root = Path(root) + self.split = split + self.mask_dir = mask_dir + self.normal_dir = normal_dir + self.test_dir = test_dir + + check_path = root + + check_path = self.root / (self.test_dir if split == Split.TEST else self.normal_dir) + + file_extension = most_common_extension(check_path) + if not file_extension: + msg = "No common file extension found in the directory." + raise ValueError(msg) + + if file_extension in SINGLE_FILE_EXTENSIONS: + self.indexer_cls = FolderClipsIndexerVideo + elif file_extension in FOLDER_IMAGE_EXTENSIONS: + self.indexer_cls = FolderClipsIndexerImgFrames + else: + msg = f"Unsupported file extension {file_extension}" + raise TypeError(msg) + + self.samples = make_folder_video_dataset( + root=self.root, + normal_dir=self.normal_dir, + test_dir=self.test_dir, + mask_dir=self.mask_dir, + split=self.split, + ) + + +class FolderVideo(AnomalibVideoDataModule): + """Folder DataModule class. + + Args: + root (Path | str): Path to the root of the dataset + normal_dir (Path | str): Path to the training videos of the dataset (.avi / .mp4 / imgages as frames). + test_dir (Path | str): Path to the testing videos of the dataset. + mask_dir (Path | str | None): Path to the masks for the training videos of the dataset (.npy/.pt/ images) + clip_length_in_frames (int, optional): Number of video frames in each clip. + 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 | str): 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. + val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained. + val_split_ratio (float): Fraction of train or test images that will be reserved for validation. + seed (int | None, optional): Seed which may be set to a fixed value for reproducibility. + """ + + def __init__( + self, + root: Path | str, + mask_dir: Path | str | None = None, + normal_dir: Path | str = "train_vid", + test_dir: Path | str = "test_vid", + clip_length_in_frames: int = 2, + 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, + val_split_mode: ValSplitMode = ValSplitMode.SAME_AS_TEST, + val_split_ratio: float = 0.5, + seed: int | None = None, + ) -> None: + super().__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, + ) + + self.task = TaskType(task) + self.root = Path(root) + self.normal_dir = Path(normal_dir) + self.test_dir = Path(test_dir) + self.mask_dir = mask_dir if mask_dir is None else Path(mask_dir) + self.clip_length_in_frames = clip_length_in_frames + self.frames_between_clips = frames_between_clips + self.target_frame = VideoTargetFrame(target_frame) + + def _setup(self, _stage: str | None = None) -> None: + self.train_data = FolderDataset( + 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, + root=self.root, + mask_dir=self.mask_dir, + normal_dir=self.normal_dir, + test_dir=self.test_dir, + split=Split.TRAIN, + ) + + self.test_data = FolderDataset( + 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, + root=self.root, + mask_dir=self.mask_dir, + normal_dir=self.normal_dir, + test_dir=self.test_dir, + split=Split.TEST, + ) + + @staticmethod + def _convert_training_videos(video_folder: Path, target_folder: Path) -> None: + """Re-code the training videos to ensure correct reading of frames by torchvision. + + The encoding of the raw video files in the ShanghaiTech dataset causes some problems when + reading the frames using pyav. To prevent this, we read the frames from the video files using opencv, + and write them to a new video file that can be parsed correctly with pyav. + + Args: + video_folder (Path): Path to the folder of training videos. + target_folder (Path): File system location where the converted videos will be stored. + """ + training_videos = sorted(video_folder.glob("*")) + for video_idx, video_path in enumerate(training_videos): + logger.info("Converting training video %s (%i/%i)...", video_path.name, video_idx + 1, len(training_videos)) + file_name = video_path.name + target_path = target_folder / file_name + convert_video(video_path, target_path, codec="XVID") diff --git a/src/anomalib/data/video/shanghaitech.py b/src/anomalib/data/video/shanghaitech.py index 6c0055dd31..822a3803e4 100644 --- a/src/anomalib/data/video/shanghaitech.py +++ b/src/anomalib/data/video/shanghaitech.py @@ -45,7 +45,7 @@ DATASET_DOWNLOAD_INFO = DownloadInfo( name="ShanghaiTech Dataset", - url="http://101.32.75.151:8181/dataset/shanghaitech.tar.gz", + url="https://101.32.75.151:8181/dataset/shanghaitech.tar.gz", hashsum="c13a827043b259ccf8493c9d9130486872992153a9d714fe229e523cd4c94116", ) diff --git a/tests/helpers/data.py b/tests/helpers/data.py index 51b683acab..f25f15313a 100644 --- a/tests/helpers/data.py +++ b/tests/helpers/data.py @@ -566,3 +566,40 @@ def _generate_dummy_shanghaitech_dataset( self.video_generator.save_image(image_filename, frame) masks_array = np.stack(masks) np.save(mask_path, masks_array) + + def _generate_dummy_foldervideo_dataset( + self, + train_dir: str = "train_vid", + test_dir: str = "test_vid", + ) -> None: + """Generate dummy folder video multiple dataset.""" + # generate training data + mask_dir = "test_labels" + path = self.dataset_root / train_dir + path.mkdir(exist_ok=True, parents=True) + num_clips = self.num_train + for clip_idx in range(num_clips): + clip_path = path / f"01_{clip_idx:03}.avi" + frames, _ = self.video_generator.generate_video(length=32, first_label=LabelName.NORMAL, p_state_switch=0) + fourcc = cv2.VideoWriter_fourcc("F", "M", "P", "4") + writer = cv2.VideoWriter(str(clip_path), fourcc, 30, self.frame_shape) + for _, frame in enumerate(frames): + writer.write(frame) + writer.release() + + # generate test data + test_path = self.dataset_root / test_dir + test_path.mkdir(exist_ok=True, parents=True) + gt_path = self.dataset_root / mask_dir + gt_path.mkdir(exist_ok=True, parents=True) + + for clip_idx in range(self.num_test): + clip_path = test_path / f"01_{clip_idx:04}" + clip_path.mkdir(exist_ok=True, parents=True) + mask_path = gt_path / f"01_{clip_idx:04}.npy" + frames, masks = self.video_generator.generate_video(length=32, p_state_switch=0.2) + for frame_idx, frame in enumerate(frames): + image_filename = clip_path / f"{frame_idx:03}.jpg" + self.video_generator.save_image(image_filename, frame) + masks_array = np.stack(masks) + np.save(mask_path, masks_array) diff --git a/tests/unit/data/video/test_foldervideo.py b/tests/unit/data/video/test_foldervideo.py new file mode 100644 index 0000000000..64692c0a10 --- /dev/null +++ b/tests/unit/data/video/test_foldervideo.py @@ -0,0 +1,44 @@ +"""Unit Tests - Folder Datamodule.""" + +# Copyright (C) 2023-2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest + +from anomalib import TaskType +from anomalib.data import FolderVideo +from tests.unit.data.base.video import _TestAnomalibVideoDatamodule + + +class TestFolderVideo(_TestAnomalibVideoDatamodule): + """FolderVideo Datamodule Unit Tests. + + All of the folder datamodule tests are placed in ``TestFolder`` class. + """ + + @pytest.fixture() + def clip_length_in_frames(self) -> int: + """Return the number of frames in each clip.""" + return 2 + + @pytest.fixture() + def datamodule(self, dataset_path: Path, task_type: TaskType, clip_length_in_frames: int) -> FolderVideo: + """Create and return a FolderVideo datamodule.""" + _datamodule = FolderVideo( + root=dataset_path / "foldervideo", + normal_dir="train_vid", + test_dir="test_vid", + mask_dir="test_labels", + 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, + ) + _datamodule.prepare_data() + _datamodule.setup() + + return _datamodule