VisA Anomaly Classifier: Complete Runnable Setup
A Mixture of Experts (MoE) image classification system that identifies industrial product categories and detects anomalies using a product classification router and 12 specialized ResNet-18 experts.
Hugging Face Model: https://huggingface.co/bipulsardar/Anomaly_Classifier
1. Project Structure
Create the following project structure:
visa-anomaly-classifier/
├── app.py
├── requirements.txt
└── README.md
Model weights are downloaded automatically from Hugging Face when the application starts.
2. Requirements
- Python 3.10 or later
- Internet access for downloading model weights
- CPU or CUDA-compatible GPU
- Approximately 600 MB for model weights, plus additional runtime memory
3. Install Dependencies
Create a virtual environment and install the dependencies.
Windows
python -m venv .venv
.venv\Scripts\activate
pip install -r requirements.txt
Linux / macOS / Google Colab
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
Create requirements.txt:
torch
torchvision
huggingface_hub
pillow
numpy
pandas
scikit-learn
tqdm
For GPU acceleration, install the PyTorch build compatible with your CUDA environment using the official selector: https://pytorch.org/get-started/locally/
4. Complete Application
Create app.py with the following code.
import argparse
import json
import time
from pathlib import Path
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from PIL import Image
from torchvision import models, transforms
from huggingface_hub import snapshot_download
from sklearn.metrics import (
accuracy_score,
precision_score,
recall_score,
f1_score,
roc_auc_score,
classification_report,
confusion_matrix
)
CATEGORIES = [
"candle", "capsules", "cashew", "chewinggum",
"fryum", "macaroni1", "macaroni2", "pcb1",
"pcb2", "pcb3", "pcb4", "pipe_fryum"
]
IMAGE_EXTENSIONS = {
".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"
}
REPOSITORY = "bipulsardar/Anomaly_Classifier"
IMAGE_SIZE = 224
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TRANSFORM = transforms.Compose([
transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
def create_model(num_classes):
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, num_classes)
return model
def load_model(path, num_classes):
checkpoint = torch.load(
path,
map_location=DEVICE,
weights_only=True
)
state_dict = checkpoint.get("model_state_dict", checkpoint)
model = create_model(num_classes)
model.load_state_dict(state_dict)
model.to(DEVICE)
model.eval()
return model
def load_models(model_dir=None, token=None):
if model_dir:
directory = Path(model_dir)
else:
directory = Path(snapshot_download(
repo_id=REPOSITORY,
repo_type="model",
token=token
))
router = load_model(directory / "router.pth", len(CATEGORIES))
experts = {
category: load_model(
directory / f"expert_{category}.pth",
2
)
for category in CATEGORIES
}
return router, experts
def predict_image(image_path, router, experts):
with Image.open(image_path) as image:
tensor = TRANSFORM(image.convert("RGB")).unsqueeze(0).to(DEVICE)
if DEVICE.type == "cuda":
torch.cuda.synchronize()
start = time.perf_counter()
with torch.inference_mode():
router_probabilities = torch.softmax(router(tensor), dim=1)
category_id = int(router_probabilities.argmax(dim=1).item())
category = CATEGORIES[category_id]
expert_probabilities = torch.softmax(
experts[category](tensor),
dim=1
)
anomaly_score = float(expert_probabilities[0, 1].item())
prediction = int(expert_probabilities.argmax(dim=1).item())
if DEVICE.type == "cuda":
torch.cuda.synchronize()
elapsed_ms = (time.perf_counter() - start) * 1000
return {
"image_path": str(image_path),
"category": category,
"category_confidence": float(
router_probabilities[0, category_id].item()
),
"prediction": "anomaly" if prediction == 1 else "normal",
"anomaly_score": anomaly_score,
"latency_ms": elapsed_ms
}
def discover_dataset(dataset_dir):
root = Path(dataset_dir)
records = []
for category in CATEGORIES:
image_root = root / category / "Data" / "Images"
for folder, label in [("Normal", 0), ("Anomaly", 1)]:
directory = image_root / folder
if not directory.exists():
continue
for path in directory.rglob("*"):
if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS:
records.append({
"image_path": str(path),
"true_category": category,
"true_label": label
})
dataframe = pd.DataFrame(records)
if dataframe.empty:
raise RuntimeError(
f"No VisA images found under {root}. "
"Expected category/Data/Images/Normal and Anomaly folders."
)
return dataframe.drop_duplicates("image_path").reset_index(drop=True)
def benchmark(dataset_dir, router, experts, output_dir):
output = Path(output_dir)
output.mkdir(parents=True, exist_ok=True)
dataframe = discover_dataset(dataset_dir)
results = []
for row in dataframe.itertuples(index=False):
try:
result = predict_image(row.image_path, router, experts)
result["true_category"] = row.true_category
result["true_label"] = int(row.true_label)
result["category_correct"] = (
result["category"] == row.true_category
)
results.append(result)
except Exception as error:
print(f"Failed: {row.image_path}: {error}")
result_df = pd.DataFrame(results)
if result_df.empty:
raise RuntimeError("No images were successfully evaluated.")
y_true = result_df["true_label"].to_numpy()
y_pred = (
result_df["prediction"].eq("anomaly").astype(int).to_numpy()
)
y_score = result_df["anomaly_score"].to_numpy()
try:
auc = float(roc_auc_score(y_true, y_score))
except ValueError:
auc = None
summary = {
"repository": REPOSITORY,
"device": str(DEVICE),
"images_evaluated": len(result_df),
"images_failed": len(dataframe) - len(result_df),
"router_accuracy": float(
result_df["category_correct"].mean()
),
"anomaly_accuracy": float(accuracy_score(y_true, y_pred)),
"anomaly_precision": float(
precision_score(y_true, y_pred, zero_division=0)
),
"anomaly_recall": float(
recall_score(y_true, y_pred, zero_division=0)
),
"anomaly_f1": float(
f1_score(y_true, y_pred, zero_division=0)
),
"anomaly_roc_auc": auc,
"mean_latency_ms": float(result_df["latency_ms"].mean()),
"median_latency_ms": float(result_df["latency_ms"].median()),
"p95_latency_ms": float(result_df["latency_ms"].quantile(0.95))
}
result_df.to_csv(output / "predictions.csv", index=False)
with open(output / "benchmark_summary.json", "w") as file:
json.dump(summary, file, indent=2)
print("\nBenchmark summary")
print(json.dumps(summary, indent=2))
print("\nClassification report")
print(classification_report(
y_true,
y_pred,
labels=[0, 1],
target_names=["Normal", "Anomaly"],
zero_division=0
))
print("\nConfusion matrix")
print(confusion_matrix(y_true, y_pred, labels=[0, 1]))
return result_df, summary
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--image", type=str)
parser.add_argument("--image-dir", type=str)
parser.add_argument("--dataset-dir", type=str)
parser.add_argument("--output-dir", default="benchmark_results")
parser.add_argument("--model-dir", type=str)
parser.add_argument("--hf-token", type=str)
args = parser.parse_args()
router, experts = load_models(args.model_dir, args.hf_token)
print(f"Device: {DEVICE}")
print(f"Loaded {len(experts)} experts.")
if args.image:
print(json.dumps(
predict_image(args.image, router, experts),
indent=2
))
return
if args.dataset_dir:
benchmark(
args.dataset_dir,
router,
experts,
args.output_dir
)
return
if args.image_dir:
image_paths = [
path for path in Path(args.image_dir).rglob("*")
if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
]
output = []
for path in image_paths:
try:
output.append(predict_image(path, router, experts))
except Exception as error:
output.append({
"image_path": str(path),
"error": str(error)
})
Path(args.output_dir).mkdir(parents=True, exist_ok=True)
pd.DataFrame(output).to_csv(
Path(args.output_dir) / "predictions.csv",
index=False
)
print(f"Processed {len(output)} images.")
return
parser.error(
"Specify --image, --image-dir, or --dataset-dir."
)
if __name__ == "__main__":
main()
5. Run Predictions
Single image
python app.py --image ./product.jpg
Batch image directory
python app.py --image-dir ./product_images --output-dir ./results
The predictions will be saved to results/predictions.csv.
Evaluate the VisA dataset
If the dataset is extracted to /content/VisA in Google Colab:
python app.py --dataset-dir /content/VisA --output-dir /content/benchmark_results
On Windows:
python app.py --dataset-dir "E:\Datasets\VisA" --output-dir ".\benchmark_results"
The benchmark outputs include per-image predictions and aggregate metrics.
6. Use in Google Colab
Run these cells in order.
!pip -q install torch torchvision huggingface_hub pillow numpy pandas scikit-learn tqdm
Upload app.py to Colab or create it using a notebook cell, then run:
!python app.py --image "/content/product.jpg"
Benchmark all VisA images:
!python app.py --dataset-dir "/content/VisA" --output-dir "/content/benchmark_results"
7. Private Hugging Face Repositories
If the repository is private, authenticate before running the application.
from huggingface_hub import login
from kaggle_secrets import UserSecretsClient
secrets = UserSecretsClient()
token = secrets.get_secret("HF_091")
login(token=token)
Alternatively, set HF_TOKEN in the environment and pass it to load_models(). Do not commit access tokens to source control.
8. Interpreting Results
The application returns:
category: predicted product category.category_confidence: router softmax score for the selected category.prediction: normal or anomaly.anomaly_score: expert softmax score for the anomaly class.latency_ms: inference time for the router and selected expert.
These scores are not necessarily calibrated probabilities. Set production thresholds using a separate validation dataset.
9. Important Evaluation Notes
The previously reported full-dataset results are preliminary because the evaluation included training images. For a reliable estimate of generalization, evaluate on held-out test images that were never used to train or select the models.
This system performs image-level anomaly classification. It does not currently produce pixel-level defect localization.
10. References
- Model repository: https://huggingface.co/bipulsardar/Anomaly_Classifier
- VisA dataset: https://github.com/amazon-science/spot-diff
- PyTorch: https://pytorch.org/
- Torchvision: https://pytorch.org/vision/
- Hugging Face Hub: https://huggingface.co/docs/huggingface_hub
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support