✨ Add a custom dataset class in anomalib.data

[mdciri]

Problem Description

How do I train a model using my own dataset? how do I create my dataset such that it is compatible with engine.fit()?
In this way, everyone will more likely use this library to train a model with using a custom dataset.

Proposed Solution

from anomalib.data import CustomDataModule
from anomalib.models import Patchcore
from anomalib.engine import Engine

# Initialize components
datamodule = CustomDataModule(???)
model = Patchcore()
engine = Engine()

# Train the model
engine.fit(datamodule=datamodule, model=model)

Alternatives Considered

No response

Additional Context

No response

[github-actions[bot]]

Thanks for submitting this issue! It has been added to our triage queue. A maintainer will review it shortly.

[rajeshgangireddy]

Hi @mdciri
You can use the Folder module to create a data module.
A simple example looks like this :

from anomalib.data import Folder
from anomalib.models import Patchcore
from anomalib.engine import Engine


dataset_folder = "path/to/dataset"

datamodule = Folder(
    name="DatasetName",  # A unique name for your dataset
    root=dataset_folder,  # Root directory containing 'normal' and 'abnormal' folders
    normal_dir="train/normal",  # Name of the folder containing normal images
    abnormal_dir = "test/abnormal",
    normal_test_dir = "test/normal",
    # explore other inputs to this
    train_batch_size=32,
    eval_batch_size=32,
    num_workers=8,
)

datamodule.setup()

model = Patchcore()
engine = Engine()
engine.fit(datamodule=datamodule, model=model)

Would that help you ?