#!/usr/bin/env python3 """ Inference script for MVTec AD using Anomalib models. Visualizes heatmap and masks of detected anomalies. Examples: python inference.py --image_path path/to/image.png python inference.py --image_path image.png --model padim python inference.py --image_path image.png --category cable """ import os import argparse from pathlib import Path import numpy as np import matplotlib.pyplot as plt from PIL import Image from anomalib.data import PredictDataset from anomalib.engine import Engine from core import ( MVTEC_CATEGORIES, DIR_RESULTS, DIR_OUTPUT, get_available_models, load_model_config, load_model, get_checkpoint_path, resize_to_match, scale_efficientad_score, ) def run_inference(image_path: str, category: str = "bottle", model_name: str = "patchcore", checkpoint_path: str = None) -> dict: """ Runs inference on a single image. Args: image_path: Path to the image category: Category for the model model_name: Name of the model to use checkpoint_path: Path to checkpoint (optional, uses default if not provided) Returns: dict with: anomaly_score, anomaly_map, pred_mask, image_path """ ckpt = checkpoint_path or str(get_checkpoint_path(category, model_name)) if not os.path.exists(ckpt): raise FileNotFoundError(f"Checkpoint not found: {ckpt}") model = load_model(model_name) dataset = PredictDataset(path=image_path) engine = Engine( default_root_dir="/tmp/anomalib_inference", callbacks=[], ) predictions = engine.predict(model=model, dataset=dataset, ckpt_path=ckpt) results = {"image_path": image_path, "category": category, "model": model_name} for batch in predictions: results["anomaly_score"] = batch.pred_score[0].cpu().item() if batch.pred_score is not None else None results["anomaly_map"] = batch.anomaly_map[0].cpu().numpy() if batch.anomaly_map is not None else None results["pred_mask"] = batch.pred_mask[0].cpu().numpy() if batch.pred_mask is not None else None break return results def visualize_results(results: dict, output_dir: Path = None) -> str: """ Visualizes heatmap and mask overlayed on original image. Args: results: Results from run_inference() output_dir: Output directory (optional) Returns: Path of saved image """ model_name = results.get("model", "unknown") base_output_dir = output_dir or DIR_OUTPUT output_dir = base_output_dir / model_name output_dir.mkdir(parents=True, exist_ok=True) image_path = results["image_path"] anomaly_score = results["anomaly_score"] anomaly_map = results["anomaly_map"] pred_mask = results.get("pred_mask") model_name = results.get("model", "unknown") original = np.array(Image.open(image_path).convert("RGB")) if anomaly_map.ndim == 3: anomaly_map = anomaly_map.squeeze(0) is_efficientad = model_name.lower() == "efficientad" amap_min, amap_max = anomaly_map.min(), anomaly_map.max() amap_range = amap_max - amap_min if is_efficientad and amap_range < 0.1: if amap_range > 1e-8: anomaly_map = (anomaly_map - amap_min) / amap_range else: anomaly_map = np.zeros_like(anomaly_map) print(f"[INFO] Applied image-level normalization (original range: {amap_min:.4f}-{amap_max:.4f})") else: anomaly_map = np.clip(anomaly_map, 0, 1) anomaly_map = resize_to_match(anomaly_map, original.shape[:2]) is_good = anomaly_score is not None and anomaly_score < 0.5 show_mask_contours = True if is_efficientad and is_good: anomaly_map = anomaly_map * 0.3 show_mask_contours = False print(f"[INFO] EfficientAD: Good image detected - scaling heatmap to low values") if pred_mask is not None: if pred_mask.ndim == 3: pred_mask = pred_mask.squeeze(0) pred_mask = resize_to_match(pred_mask, original.shape[:2]) show_fourth_panel = pred_mask is not None or (is_efficientad and is_good) num_cols = 4 if show_fourth_panel else 3 fig, axes = plt.subplots(1, num_cols, figsize=(5 * num_cols, 5), facecolor='white') for ax in axes: ax.set_facecolor('white') axes[0].imshow(original) axes[0].set_title("Original") axes[0].axis("off") if is_efficientad: anomaly_map_masked = np.ma.masked_where(anomaly_map == 0, anomaly_map) cmap = plt.cm.jet cmap.set_bad(color='none') else: anomaly_map_masked = anomaly_map cmap = plt.cm.jet im = axes[1].imshow(anomaly_map_masked, cmap=cmap, vmin=0, vmax=1, aspect='auto') axes[1].set_title(f"Anomaly Heatmap ({model_name})") axes[1].axis("off") plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04) axes[2].imshow(original, aspect='auto') axes[2].imshow(anomaly_map_masked, cmap=cmap, alpha=0.5, vmin=0, vmax=1, aspect='auto') axes[2].set_title("Overlay") axes[2].axis("off") if show_fourth_panel: axes[3].imshow(original) if show_mask_contours and pred_mask is not None: axes[3].contour(pred_mask, levels=[0.5], colors="red", linewidths=2) axes[3].set_title("Predicted Mask") axes[3].axis("off") if anomaly_score is not None: if is_efficientad: scaled_score = scale_efficientad_score(anomaly_score) score_str = f"{scaled_score:.4f}" else: score_str = f"{anomaly_score:.4f}" else: score_str = "N/A" plt.suptitle(f"Model: {model_name} | Anomaly Score: {score_str}", fontsize=14) plt.tight_layout() image_name = Path(image_path).stem output_path = output_dir / f"inference_{image_name}.png" plt.savefig(output_path, dpi=150, bbox_inches="tight") plt.show() plt.close(fig) return str(output_path) def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser( description="Inference script for MVTec AD", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: python inference.py --image_path test.png python inference.py --image_path test.png --model padim python inference.py --image_path test.png --category cable """ ) parser.add_argument( "--image_path", type=str, required=True, help="Path to image to analyze" ) parser.add_argument( "--category", type=str, default="bottle", choices=MVTEC_CATEGORIES, help="Category of the model (default: bottle)" ) parser.add_argument( "--model", type=str, default="patchcore", choices=get_available_models(), help="Model to use (default: patchcore)" ) parser.add_argument( "--checkpoint", type=str, default=None, help="Checkpoint path (optional, uses default for category/model)" ) parser.add_argument( "--output_dir", type=str, default=None, help="Output directory (optional)" ) return parser.parse_args() def main(): args = parse_args() if not os.path.exists(args.image_path): raise FileNotFoundError(f"Image not found: {args.image_path}") print(f"Inference on: {args.image_path}") print(f"Model: {args.model} | Category: {args.category}") results = run_inference( image_path=args.image_path, category=args.category, model_name=args.model, checkpoint_path=args.checkpoint ) score = results["anomaly_score"] print(f"\n{'='*50}") print(f"ANOMALY SCORE: {score:.4f}" if score else "ANOMALY SCORE: N/A") print(f"{'='*50}\n") output_dir = Path(args.output_dir) if args.output_dir else None output_path = visualize_results(results, output_dir) print(f"Saved: {output_path}") return results if __name__ == "__main__": main()