39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from anomalib.data import MVTec
|
|
from anomalib.engine import Engine
|
|
from anomalib.models import EfficientAd
|
|
|
|
def main():
|
|
# 1. Initialize the MVTec datamodule (defaults to downloading/using the 'bottle' category or similar)
|
|
datamodule = MVTec(
|
|
category="bottle",
|
|
train_batch_size=16,
|
|
eval_batch_size=16,
|
|
)
|
|
|
|
# 2. Choose an anomaly detection model
|
|
model = EfficientAd()
|
|
|
|
# 3. Initialize the PyTorch Lightning based Engine
|
|
# (Setting max_epochs low here just for a quick demonstration)
|
|
engine = Engine(
|
|
max_epochs=5,
|
|
default_root_dir="results"
|
|
)
|
|
|
|
print("Starting training...")
|
|
# 4. Train the model
|
|
engine.fit(datamodule=datamodule, model=model)
|
|
|
|
print("Running prediction/inference on test images...")
|
|
# 5. Run prediction using the best checkpoint generated during training
|
|
predictions = engine.predict(
|
|
datamodule=datamodule,
|
|
model=model,
|
|
ckpt_path=engine.trainer.checkpoint_callback.best_model_path,
|
|
)
|
|
|
|
print(f"Predictions complete. Results generated for {len(predictions)} batches.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|