Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| from torchvision import transforms | |
| from PIL import Image | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| # PatchCore import for your config (update if needed for your Anomalib version) | |
| from anomalib.models.image.patchcore import Patchcore | |
| # Config from your YAML | |
| CKPT_PATH = "results/Patchcore/transformers/v2/weights/lightning/model.ckpt" | |
| IMAGE_SIZE = (256, 256) # Update if you used a different size in training | |
| # Inference directory (change as needed) | |
| INFER_DIR = "./dataset/test/faulty" | |
| # Output directory for masks | |
| OUT_MASK_DIR = "inference_masks" | |
| os.makedirs(OUT_MASK_DIR, exist_ok=True) | |
| # Load model | |
| model = Patchcore.load_from_checkpoint(CKPT_PATH) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = model.to(device) | |
| model.eval() | |
| # Transforms (should match your training config) | |
| transform = transforms.Compose([ | |
| transforms.Resize(IMAGE_SIZE), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| # Inference loop | |
| for fname in sorted(os.listdir(INFER_DIR)): | |
| if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')): | |
| continue | |
| fpath = os.path.join(INFER_DIR, fname) | |
| img = Image.open(fpath).convert("RGB") | |
| img_tensor = transform(img).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| output = model(img_tensor) | |
| # PatchCore returns (anomaly_score, anomaly_map, ...) | |
| if hasattr(output, 'anomaly_map'): | |
| anomaly_map = output.anomaly_map.squeeze().cpu().numpy() | |
| elif isinstance(output, (tuple, list)) and len(output) > 1: | |
| anomaly_map = output[1].squeeze().cpu().numpy() | |
| else: | |
| anomaly_map = None | |
| # Save mask as PNG | |
| if anomaly_map is not None: | |
| # Normalize to 0-255 for visualization | |
| norm_map = (255 * (anomaly_map - anomaly_map.min()) / (np.ptp(anomaly_map) + 1e-8)).astype(np.uint8) | |
| mask_img = Image.fromarray(norm_map) | |
| mask_img.save(os.path.join(OUT_MASK_DIR, f"{os.path.splitext(fname)[0]}_mask.png")) | |
| print(f"Saved mask for {fname}") | |
| else: | |
| print(f"No mask generated for {fname}") | |
| print(f"All masks saved to {OUT_MASK_DIR}") | |