first commit

train.py

@@ -0,0 +1,38 @@
+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()