""" MVTec Anomaly Detection Demo — Gradio app (no FastAPI required). Run: python gradio_app.py Environment variables (all optional): ANOMAVISION_MODEL_DATA_PATH path that contains the model file (default: "distributions/anomav_exp") ANOMAVISION_MODEL_FILE model filename (default: "padim_model.onnx") ANOMAVISION_DEVICE "auto" | "cpu" | "cuda" (default: "auto") ANOMAVISION_THRESHOLD float anomaly threshold (default: 13.0) ANOMAVISION_VIZ_PADDING int, boundary-frame padding (default: 40) ANOMAVISION_VIZ_ALPHA float, heatmap blend alpha (default: 0.5) ANOMAVISION_VIZ_COLOR R,G,B highlight color (default: 128,0,128) SAMPLE_IMAGES_DIR directory with sample images (default: "sample_images") """ import io import os import time import glob import base64 from pathlib import Path from typing import Optional, Tuple import gradio as gr import numpy as np import torch from PIL import Image # HF Hub model download (only used when running on HF Spaces) try: from huggingface_hub import hf_hub_download HF_HUB_AVAILABLE = True except ImportError: HF_HUB_AVAILABLE = False # ── lazy import so the app still starts even if anomavision isn't installed ── try: import anomavision from anomavision.general import determine_device from anomavision.inference.model.wrapper import ModelWrapper from anomavision.inference.modelType import ModelType ANOMAVISION_AVAILABLE = True except ImportError: ANOMAVISION_AVAILABLE = False print("WARNING: anomavision not found – running in DEMO mode (random scores).") # ───────────────────────────────────────────────────────────────────────────── # Config # ───────────────────────────────────────────────────────────────────────────── _SCRIPT_DIR = Path(__file__).resolve().parent MODEL_DATA_PATH = os.getenv("ANOMAVISION_MODEL_DATA_PATH", str(_SCRIPT_DIR / "distributions" / "anomav_exp")) MODEL_FILE = os.getenv("ANOMAVISION_MODEL_FILE", "padim_model.onnx") DEVICE_ENV = os.getenv("ANOMAVISION_DEVICE", "auto") THRESHOLD_DEFAULT = float(os.getenv("ANOMAVISION_THRESHOLD", "13.0")) VIZ_PADDING = int(os.getenv("ANOMAVISION_VIZ_PADDING", "40")) VIZ_ALPHA = float(os.getenv("ANOMAVISION_VIZ_ALPHA", "0.5")) VIZ_COLOR = tuple(map(int, os.getenv("ANOMAVISION_VIZ_COLOR", "128,0,128").split(","))) SAMPLE_DIR = os.getenv("SAMPLE_IMAGES_DIR", str(_SCRIPT_DIR / "sample_images")) # ───────────────────────────────────────────────────────────────────────────── # Model — loaded once at startup # ───────────────────────────────────────────────────────────────────────────── _model: Optional["ModelWrapper"] = None _model_type = None _device_str: str = "cpu" def _load_model() -> str: """Load the model and return a status message.""" global _model, _model_type, _device_str if not ANOMAVISION_AVAILABLE: return "⚠️ anomavision not installed — running in demo mode." model_path = os.path.realpath(os.path.join(MODEL_DATA_PATH, MODEL_FILE)) # If model not found locally, try downloading from HF Hub if not os.path.exists(model_path) and HF_HUB_AVAILABLE: hf_repo = os.getenv("ANOMAVISION_HF_MODEL_REPO", "") if hf_repo: try: print(f"Downloading model from HF Hub: {hf_repo}/{MODEL_FILE}") os.makedirs(MODEL_DATA_PATH, exist_ok=True) model_path = hf_hub_download( repo_id=hf_repo, filename=MODEL_FILE, local_dir=MODEL_DATA_PATH, ) print(f"Model downloaded to: {model_path}") except Exception as e: return f"⚠️ HF Hub download failed: {e} — running in demo mode." else: return f"⚠️ Model not found at {model_path} and ANOMAVISION_HF_MODEL_REPO not set — running in demo mode." if not os.path.exists(model_path): return f"⚠️ Model not found at {model_path} — running in demo mode." try: _device_str = determine_device(DEVICE_ENV) _model_type = ModelType.from_extension(model_path) _model = ModelWrapper(model_path, _device_str) # Optional warmup try: dummy = torch.zeros((1, 3, 224, 224), dtype=torch.float32, device=_device_str) _model.warmup(batch=dummy, runs=1) except Exception: pass return f"✅ Model loaded: {Path(model_path).name} ({_model_type.value}) on {_device_str}" except Exception as e: return f"⚠️ Model load failed: {e} — running in demo mode." _startup_message = _load_model() print(_startup_message) # ───────────────────────────────────────────────────────────────────────────── # Inference helpers # ───────────────────────────────────────────────────────────────────────────── def _pil_to_np(image: Image.Image) -> np.ndarray: return np.array(image.convert("RGB")) def _demo_predict(image_np: np.ndarray): """Return fake results when the real model isn't available.""" h, w = image_np.shape[:2] score = float(np.random.uniform(5, 25)) # random heatmap heatmap_np = np.random.rand(h, w).astype(np.float32) return score, heatmap_np def _real_predict(image_np: np.ndarray, threshold: float): """Run anomavision inference and return (score, score_map_np, boundary_np, heatmap_np, highlighted_np).""" device = torch.device(_device_str) batch = anomavision.to_batch([image_np], anomavision.standard_image_transform, device) if _device_str == "cuda": batch = batch.half() with torch.no_grad(): image_scores, score_maps = _model.predict(batch) # Update threshold temporarily for visualizations _orig_threshold = getattr(anomavision, "_threshold", None) score_map_cls = anomavision.classification(score_maps, threshold) image_cls = anomavision.classification(image_scores, threshold) test_images = np.array([image_np]) boundary_images = anomavision.visualization.framed_boundary_images( test_images, score_map_cls, image_cls, padding=VIZ_PADDING ) heatmap_images = anomavision.visualization.heatmap_images( test_images, score_maps, alpha=VIZ_ALPHA ) highlighted_images = anomavision.visualization.highlighted_images( [image_np], score_map_cls, color=VIZ_COLOR ) sm0 = score_maps[0] if isinstance(sm0, np.ndarray): score_map_np = sm0 elif hasattr(sm0, "cpu"): score_map_np = sm0.cpu().float().numpy() else: score_map_np = np.array(sm0) return ( float(image_scores[0]), score_map_np, boundary_images[0], heatmap_images[0], highlighted_images[0], ) def _np_to_pil(arr: np.ndarray, size: Optional[Tuple[int, int]] = None) -> Image.Image: if arr is None: return None if arr.dtype != np.uint8: if arr.max() <= 1.0: arr = (arr * 255).astype(np.uint8) else: arr = np.clip(arr, 0, 255).astype(np.uint8) img = Image.fromarray(arr) if size: img = img.resize(size, Image.BILINEAR) return img # ───────────────────────────────────────────────────────────────────────────── # Sample images # ───────────────────────────────────────────────────────────────────────────── SUPPORTED_EXT = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} def _collect_samples() -> list: """ Collect sample images from SAMPLE_DIR. Expected layout (mirrors MVTec): sample_images/ bottle/broken_large/000.png bottle/good/001.png cable/bent_wire/000.png … Falls back to any image recursively found in SAMPLE_DIR. Returns list of (display_label, abs_path). """ samples = [] base = Path(SAMPLE_DIR) if not base.exists(): return samples for p in sorted(base.rglob("*")): if p.suffix.lower() in SUPPORTED_EXT: # Build label: "category/defect_type" rel = p.relative_to(base) parts = rel.parts if len(parts) >= 3: label = f"{parts[0]}/{parts[1]}" elif len(parts) == 2: label = f"{parts[0]}/{p.stem}" else: label = p.stem samples.append((label, str(p))) return samples SAMPLES = _collect_samples() # ───────────────────────────────────────────────────────────────────────────── # Main inference function (called by Gradio) # ───────────────────────────────────────────────────────────────────────────── def run_inference( image: Optional[Image.Image], threshold: float, resize_w: int, resize_h: int, include_viz: bool, ) -> Tuple: if image is None: return "❌ Please upload or select an image.", None, None, None, None resize = (int(resize_w), int(resize_h)) image_np = _pil_to_np(image) t0 = time.time() if _model is not None and ANOMAVISION_AVAILABLE: try: score, score_map_np, boundary_np, heatmap_np, highlighted_np = _real_predict( image_np, threshold ) is_anomaly = score >= threshold original_pil = image.resize(resize, Image.BILINEAR) heatmap_pil = _np_to_pil(heatmap_np, resize) if include_viz else None boundary_pil = _np_to_pil(boundary_np, resize) if include_viz else None highlighted_pil = _np_to_pil(highlighted_np, resize) if include_viz else None except Exception as e: return f"⚠️ Inference error: {e}", None, None, None, None else: # Demo mode score, heatmap_raw = _demo_predict(image_np) is_anomaly = score >= threshold original_pil = image.resize(resize, Image.BILINEAR) if include_viz: # Generate a fake colorised heatmap import matplotlib.cm as cm heatmap_norm = heatmap_raw / heatmap_raw.max() cmap = cm.get_cmap("jet") heatmap_rgba = (cmap(heatmap_norm) * 255).astype(np.uint8) heatmap_rgb = heatmap_rgba[:, :, :3] blend = (0.5 * image_np + 0.5 * heatmap_rgb).astype(np.uint8) heatmap_pil = _np_to_pil(heatmap_rgb, resize) boundary_pil = _np_to_pil(blend, resize) highlighted_pil = _np_to_pil(image_np, resize) else: heatmap_pil = boundary_pil = highlighted_pil = None elapsed = time.time() - t0 label = "🚨 ANOMALY DETECTED" if is_anomaly else "✅ NORMAL" status = f"Model: {Path(MODEL_FILE).stem} | Score: {score:.4f} | {label}" detail = f"Threshold: {threshold:.2f} | Inference time: {elapsed:.2f}s" return f"{status}\n{detail}", original_pil, heatmap_pil, boundary_pil, highlighted_pil # ───────────────────────────────────────────────────────────────────────────── # Build Gradio UI # ───────────────────────────────────────────────────────────────────────────── _ACCENT = "#f97316" # orange _BG = "#ffffff" _SURFACE = "#f8fafc" _BORDER = "#e2e8f0" _TEXT = "#0f172a" _MUTED = "#64748b" custom_css = f""" /* ── Global ── */ body, .gradio-container {{ background: {_BG} !important; color: {_TEXT} !important; font-family: 'DM Sans', 'Segoe UI', sans-serif !important; }} /* Force light mode */ :root {{ color-scheme: light; }} /* header */ .app-header {{ padding: 2rem 2.5rem 1rem; border-bottom: 1px solid {_BORDER}; margin-bottom: 1.5rem; }} .app-header h1 {{ font-size: 1.7rem; font-weight: 700; letter-spacing: -0.02em; margin: 0; color: {_TEXT}; }} .app-header p {{ margin: 0.25rem 0 0; color: {_MUTED}; font-size: 0.9rem; }} .status-pill {{ display: inline-block; padding: 0.2rem 0.75rem; border-radius: 99px; background: #dcfce7; color: #16a34a; border: 1px solid #86efac; font-size: 0.75rem; font-weight: 600; margin-left: 0.75rem; vertical-align: middle; }} /* Tabs */ .tab-nav button {{ background: transparent !important; color: {_MUTED} !important; border-bottom: 2px solid transparent !important; border-radius: 0 !important; font-weight: 500 !important; padding: 0.6rem 1.1rem !important; transition: color .2s, border-color .2s !important; }} .tab-nav button.selected {{ color: {_ACCENT} !important; border-bottom-color: {_ACCENT} !important; }} /* Panels */ .panel-box {{ background: {_SURFACE}; border: 1px solid {_BORDER}; border-radius: 10px; padding: 1rem 1.25rem; }} .panel-label {{ font-size: 0.78rem; font-weight: 600; letter-spacing: 0.05em; text-transform: uppercase; color: {_MUTED}; margin-bottom: 0.6rem; display: flex; align-items: center; gap: 0.4rem; }} .panel-label span.dot {{ width: 8px; height: 8px; border-radius: 50%; background: {_ACCENT}; display: inline-block; }} /* Result header */ .result-header {{ background: {_SURFACE}; border: 1px solid {_BORDER}; border-radius: 8px; padding: 0.75rem 1rem; font-family: 'JetBrains Mono', monospace; font-size: 0.82rem; white-space: pre-wrap; color: {_TEXT}; min-height: 3.5rem; }} /* Analyze button */ .btn-analyze {{ background: {_ACCENT} !important; color: #fff !important; font-weight: 700 !important; font-size: 1rem !important; border-radius: 8px !important; border: none !important; padding: 0.75rem !important; cursor: pointer !important; width: 100% !important; transition: opacity .2s !important; }} .btn-analyze:hover {{ opacity: 0.88 !important; }} /* Sample gallery */ .sample-gallery {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(110px, 1fr)); gap: 8px; max-height: 280px; overflow-y: auto; padding: 4px; }} .sample-thumb {{ cursor: pointer; border-radius: 8px; overflow: hidden; border: 2px solid transparent; transition: border-color .15s, transform .15s; background: {_BORDER}; position: relative; }} .sample-thumb:hover {{ border-color: {_ACCENT}; transform: scale(1.03); }} .sample-thumb img {{ width: 100%; aspect-ratio: 1; object-fit: cover; display: block; }} .sample-thumb .thumb-label {{ font-size: 0.62rem; color: {_MUTED}; text-align: center; padding: 3px 4px 4px; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; background: {_SURFACE}; }} /* Image panels */ .image-grid {{ display: grid; grid-template-columns: repeat(4, 1fr); gap: 10px; }} .img-card {{ background: {_SURFACE}; border: 1px solid {_BORDER}; border-radius: 10px; overflow: hidden; }} .img-card-title {{ font-size: 0.72rem; font-weight: 600; letter-spacing: 0.04em; text-transform: uppercase; color: {_MUTED}; text-align: center; padding: 6px 8px 0; }} /* Sliders, inputs */ input[type=range] {{ accent-color: {_ACCENT} !important; }} .gr-input, .gr-box, .gr-form, textarea, input[type=number] {{ background: #ffffff !important; border-color: {_BORDER} !important; color: {_TEXT} !important; border-radius: 6px !important; }} label span {{ color: {_MUTED} !important; font-size: 0.83rem !important; }} /* Scrollbar */ ::-webkit-scrollbar {{ width: 5px; height: 5px; }} ::-webkit-scrollbar-track {{ background: {_BG}; }} ::-webkit-scrollbar-thumb {{ background: {_BORDER}; border-radius: 3px; }} """ def _gallery_data() -> list: """Return list of (file_path, caption) for gr.Gallery — paths load faster than PIL objects.""" items = [] for label, path in SAMPLES: if os.path.exists(path): items.append((path, label)) if not items: print(f"[WARNING] No sample images found in: {SAMPLE_DIR}") else: print(f"[INFO] Loaded {len(items)} sample images from {SAMPLE_DIR}") return items def on_gallery_select(evt: gr.SelectData) -> Optional[Image.Image]: """Called when user clicks a thumbnail in the Gallery.""" try: idx = evt.index if isinstance(idx, (list, tuple)): idx = idx[0] idx = int(idx) if 0 <= idx < len(SAMPLES): _, path = SAMPLES[idx] return Image.open(path).convert("RGB") except Exception as e: print(f"Gallery select error: {e}") return None def on_gallery_select_and_run(evt: gr.SelectData, threshold, resize_w, resize_h, viz_check): """Load sample image and immediately run inference.""" img = on_gallery_select(evt) if img is None: return None, "❌ Could not load sample image.", None, None, None, None status, orig, heat, overlay, mask = run_inference(img, threshold, resize_w, resize_h, viz_check) return img, status, orig, heat, overlay, mask # ───────────────────────────────────────────────────────────────────────────── # Gradio app # ───────────────────────────────────────────────────────────────────────────── with gr.Blocks(title="MVTec Anomaly Detection") as demo: # ── Header ────────────────────────────────────────────────────────────── gr.HTML(f"""

🔍 MVTec Anomaly Detection Demo ● Running

Upload an image or pick a sample below — get the heatmap, overlay & predicted mask instantly.

{_startup_message}

""") # ── Tabs ───────────────────────────────────────────────────────────────── with gr.Tabs(): # ── Tab 1: Upload Image ─────────────────────────────────────────────── with gr.Tab("📤 Upload Image"): with gr.Row(equal_height=False): # ── Left column: controls ──────────────────────────────────── with gr.Column(scale=1, min_width=300): gr.HTML('
Input
') input_img = gr.Image( type="pil", label="Upload Image", show_label=False, height=280, ) with gr.Row(): model_dd = gr.Dropdown( choices=[Path(MODEL_FILE).stem], value=Path(MODEL_FILE).stem, label="Model", scale=1, ) category_dd = gr.Dropdown( choices=["bottle", "cable", "carpet", "grid", "hazelnut", "leather", "metal_nut", "pill", "screw", "tile", "toothbrush", "transistor", "wood", "zipper", "other"], value="bottle", label="Category", scale=1, ) threshold = gr.Slider( 0.1, 50.0, THRESHOLD_DEFAULT, step=0.1, label="Threshold" ) with gr.Row(): resize_w = gr.Number(value=224, label="Width", minimum=32, maximum=2048, precision=0) resize_h = gr.Number(value=224, label="Height", minimum=32, maximum=2048, precision=0) viz_check = gr.Checkbox(value=True, label="Generate Visualizations") analyze_btn = gr.Button( "🔍 Analyze Image", elem_classes=["btn-analyze"], variant="primary", ) # ── Sample images ──────────────────────────────────────── gr.HTML('
Sample Images (click to select)
') sample_gallery = gr.Gallery( value=_gallery_data(), label="", show_label=False, columns=4, rows=2, height=260, object_fit="cover", allow_preview=False, ) # ── Right column: results ───────────────────────────────────── with gr.Column(scale=2): gr.HTML('
Results
') result_text = gr.Textbox( label="", lines=2, show_label=False, elem_classes=["result-header"], placeholder="Run inference to see results…", ) with gr.Row(): out_original = gr.Image(label="Original", type="pil") out_heatmap = gr.Image(label="Anomaly Heatmap", type="pil") out_overlay = gr.Image(label="Overlay", type="pil") out_mask = gr.Image(label="Predicted Mask", type="pil") # Wire up "Analyze" button analyze_btn.click( fn=run_inference, inputs=[input_img, threshold, resize_w, resize_h, viz_check], outputs=[result_text, out_original, out_heatmap, out_overlay, out_mask], ) # Wire up sample gallery click → load image into input AND run inference sample_gallery.select( fn=on_gallery_select_and_run, inputs=[threshold, resize_w, resize_h, viz_check], outputs=[input_img, result_text, out_original, out_heatmap, out_overlay, out_mask], ) # ── Tab 2: Draw Defects ─────────────────────────────────────────────── with gr.Tab("🎨 Draw Defects"): gr.Markdown(""" ## Draw Artificial Defects 1. Upload a **GOOD** (normal) image 2. Use the brush tool to paint artificial defects 3. Click Analyze to see if the model detects them *(Sketch editor — available in Gradio ≥ 4.x)* """) with gr.Row(): with gr.Column(): sketch_img = gr.ImageEditor( type="pil", label="Draw Defects Here", brush=gr.Brush(colors=["#ff0000", "#ffff00", "#ffffff"], default_size=8), ) sketch_threshold = gr.Slider(0.1, 50.0, THRESHOLD_DEFAULT, step=0.1, label="Threshold") sketch_btn = gr.Button("🔍 Analyze Drawn Image", variant="primary") with gr.Column(): sketch_result = gr.Textbox(label="Result", lines=2) sketch_heat = gr.Image(label="Heatmap", type="pil") sketch_overlay = gr.Image(label="Overlay", type="pil") def run_sketch(editor_val, thr): if editor_val is None: return "Please draw on the image first.", None, None img = editor_val.get("composite") if isinstance(editor_val, dict) else editor_val if img is None: return "Please draw on the image first.", None, None status, orig, heat, boundary, _ = run_inference(img, thr, 224, 224, True) return status, heat, boundary sketch_btn.click( fn=run_sketch, inputs=[sketch_img, sketch_threshold], outputs=[sketch_result, sketch_heat, sketch_overlay], ) # ── Tab 3: Compare Models ───────────────────────────────────────────── with gr.Tab("⚖️ Compare Models"): gr.Markdown("## Compare Models\n*Coming soon — run two models side-by-side.*") # ── Tab 4: Learn About Models ───────────────────────────────────────── with gr.Tab("📚 Learn About Models"): gr.Markdown(""" ## Understanding Anomaly Detection Models ### 🧩 PatchCore **Approach:** Memory Bank + K-Nearest Neighbors PatchCore stores a coreset of normal patch features extracted by a pretrained CNN (e.g. WideResNet-50). At test time it computes each patch's distance to its nearest neighbour in the memory bank — high distance → anomaly. | Strength | Weakness | |----------|----------| | Very high accuracy on textures | Memory grows with dataset | | No training needed | Inference speed depends on bank size | | Works with few normal samples | | --- ### 📊 PaDiM **Approach:** Multivariate Gaussian per Patch Position PaDiM fits a multivariate Gaussian (mean + covariance) at every spatial location using multi-scale CNN features from normal images. Anomaly score = Mahalanobis distance from the expected distribution. | Strength | Weakness | |----------|----------| | Memory-efficient (only statistics stored) | Covariance matrices required | | Good cross-defect generalisation | Quality depends on backbone | --- ### ⚡ FastFlow **Approach:** Normalising Flows on CNN features FastFlow trains a normalising flow on top of frozen CNN features. It learns the density of normal features; low-likelihood patches are flagged as anomalous. | Strength | Weakness | |----------|----------| | Fastest inference | Requires training a flow network | | State-of-art on structured objects | More complex to train | """) # ── Tab 5: Model Metrics ────────────────────────────────────────────── with gr.Tab("📈 Model Metrics"): gr.Markdown(""" ## Training Metrics & Performance ### 🏆 Best Performers | Metric | Best Model | Category | Value | |--------|------------|----------|-------| | Image AUROC | Fastflow | bottle | 1.0000 | | Pixel AUROC | Patchcore | carpet | 0.9907 | | F1 Score | Fastflow | bottle | 0.9920 | > Metrics are measured on the MVTec AD benchmark test sets. """) # On HF Spaces the app is launched by the platform — __main__ block is still used locally demo.launch( server_name="0.0.0.0", server_port=7860, share=False, show_error=True, theme=gr.themes.Default(), css=custom_css, )