Spaces:
Sleeping
Sleeping
File size: 1,974 Bytes
9cf599c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import os
import torch
from torchvision import transforms
from PIL import Image
from anomalib.models.image.patchcore import Patchcore
import pandas as pd
# Path to your trained checkpoint
CKPT_PATH = "results/Patchcore/transformers/v2/weights/lightning/model.ckpt"
# Folder with images to score
INPUT_DIR = "./dataset/test/faulty" # Change as needed
# Output CSV for scores
OUTPUT_CSV = "patchcore_batch_scores.csv"
# Load PatchCore model
model = Patchcore.load_from_checkpoint(CKPT_PATH)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()
# Define image transforms (should match your training transforms)
transform = transforms.Compose([
transforms.Resize((256, 256)), # Change to your image_size if different
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def load_image(path):
img = Image.open(path).convert("RGB")
return transform(img)
# Score all images in the input directory
results = []
for fname in sorted(os.listdir(INPUT_DIR)):
if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):
continue
fpath = os.path.join(INPUT_DIR, fname)
img_tensor = load_image(fpath).unsqueeze(0) # Add batch dim
with torch.no_grad():
output = model(img_tensor.to(device)) # Move to device
# PatchCore returns a namedtuple with anomaly_score
if hasattr(output, 'anomaly_score'):
score = output.anomaly_score.item()
elif isinstance(output, (tuple, list)):
score = float(output[0])
else:
raise RuntimeError("Unknown PatchCore output type: {}".format(type(output)))
results.append({"image": fname, "anomaly_score": score})
print(f"{fname}: {score:.4f}")
# Save results to CSV
pd.DataFrame(results).to_csv(OUTPUT_CSV, index=False)
print(f"Saved batch scores to {OUTPUT_CSV}")
|