initial commit

.gitignore

@@ -0,0 +1,69 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual Environment
+venv/
+ENV/
+env/
+.venv
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# Project specific
+models/*  # Cached models are large, keep them out of git
+datasets/
+*.tar.xz
+*.zip
+*.ckpt
+*.pth
+*.onnx
+
+# Environment variables
+.env
+.env.local
+.env.*.local
+
+# Logs
+*.log
+logs/
+
+# Temporary files
+tmp/
+temp/
+saved_pillow_image.png
+
+# Jupyter
+.ipynb_checkpoints/
+*.ipynb
+
+# OS
+Thumbs.db
+
+# Gradio
+gradio_cookie_*.json

README.md

@@ -1,5 +1,5 @@
 ---
-title: Iad Explainable
+title: Industrial Anomaly Detection & Explainability System
 emoji: 🌖
 colorFrom: red
 colorTo: indigo
@@ -8,7 +8,7 @@ sdk_version: 6.10.0
 app_file: app.py
 pinned: false
 license: mit
-short_description: Multi-model anomaly detection (PatchCore + EfficientAD) with
+short_description: Multi-model anomaly detection (PatchCore + EfficientAD) with Explainable AI using Gemini
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,290 @@
+"""
+Main application entry point with Gradio UI.
+
+This module orchestrates the anomaly detection pipeline by integrating
+all core modules and providing a user-friendly web interface.
+"""
+
+import gradio as gr
+import numpy as np
+import cv2
+
+from config import THRESHOLDS, MVTEC_CATEGORIES, IMAGE_SIZE, THRESHOLD_MULTIPLIER
+from models.model_loader import load_model
+from core.inference import run_inference
+from core.postprocess import postprocess
+from core.visualization import create_visuals
+from core.explain import get_explanation, init_gemini_client
+
+
+def get_threshold(model_name: str, category: str) -> float:
+    """
+    Get the threshold for a specific model and category, applying the multiplier.
+
+    Args:
+        model_name: Name of the model
+        category: MVTec AD category
+
+    Returns:
+        Adjusted threshold value
+    """
+    base_thresh = THRESHOLDS[model_name][category]
+    return base_thresh * THRESHOLD_MULTIPLIER
+
+
+def get_status(score: float, model_name: str, category: str) -> str:
+    """
+    Determine the anomaly status based on score and threshold.
+
+    Args:
+        score: Anomaly score
+        model_name: Name of the model
+        category: MVTec AD category
+
+    Returns:
+        Status string with emoji
+    """
+    threshold = get_threshold(model_name, category)
+
+    if score < threshold:
+        return "🟢 Normal"
+    elif score < threshold + 0.1:
+        return "🟡 Slight Deviation"
+    else:
+        return "🔴 Strong Anomaly"
+
+
+def is_valid_anomaly(score: float, model_name: str, category: str) -> bool:
+    """
+    Check if the score indicates a valid anomaly.
+
+    Args:
+        score: Anomaly score
+        model_name: Name of the model
+        category: MVTec AD category
+
+    Returns:
+        True if score exceeds threshold
+    """
+    threshold = get_threshold(model_name, category)
+    return score > threshold
+
+
+def scale_efficientad_score(score: float) -> float:
+    """
+    Scale EfficientAD score for better visualization and display.
+
+    Args:
+        score: Raw EfficientAD score
+
+    Returns:
+        Scaled score
+    """
+    if score < 0.5:
+        return (score * 2) ** 2 / 4
+    else:
+        k = 500
+        return 1 / (1 + np.exp(-k * (score - 0.5)))
+
+
+def detect(image, model_name: str, category: str, gemini_client):
+    """
+    Main detection function that runs the full anomaly detection pipeline.
+
+    Args:
+        image: Input image (PIL Image or numpy array)
+        model_name: Selected model name
+        category: Selected MVTec AD category
+        gemini_client: Initialized Gemini client for explanations
+
+    Returns:
+        Tuple of visualization outputs and metadata
+    """
+    if image is None:
+        return None, None, None, None, "", "", "", None
+
+    # Convert PIL Image to numpy array
+    image_np = np.array(image)
+
+    # Load model
+    model = load_model(model_name, category)
+
+    # Run inference to get raw outputs
+    heatmap, pred_mask_raw, score = run_inference(model, image_np, model_name, category)
+
+    # Determine if it's an anomaly
+    is_anomaly = is_valid_anomaly(score, model_name, category)
+
+    # Postprocess to get final mask and bounding boxes
+    # Note: We pass the resized image for postprocessing
+    img_resized = cv2.resize(image_np, (IMAGE_SIZE, IMAGE_SIZE))
+    final_mask, bboxes, heatmap_vis = postprocess(heatmap, img_resized, model_name, is_anomaly)
+
+    # Create visualizations
+    original_vis, heatmap_color, overlay, mask_vis = create_visuals(
+        image_np, heatmap_vis, final_mask, bboxes, model_name
+    )
+
+    # Get threshold and status
+    threshold = get_threshold(model_name, category)
+    status = get_status(score, model_name, category)
+
+    # Scale score for display if using EfficientAD
+    if model_name == "efficientad":
+        display_score = scale_efficientad_score(score)
+    else:
+        display_score = score
+
+    # Store state for explanation
+    state = {
+        "image": image_np,
+        "bboxes": bboxes,
+        "score": score,
+        "category": category,
+        "gemini_client": gemini_client
+    }
+
+    return (
+        original_vis,
+        heatmap_color,
+        overlay,
+        mask_vis,
+        f"{display_score:.4f}",
+        f"{threshold:.4f}",
+        status,
+        state
+    )
+
+
+def explain(state):
+    """
+    Generate an explanation for the detected anomaly.
+
+    Args:
+        state: State dictionary containing image, bboxes, score, category, and gemini_client
+
+    Returns:
+        Explanation text
+    """
+    if state is None:
+        return "Run detection first."
+
+    gemini_client = state.get("gemini_client")
+    if gemini_client is None:
+        return "Gemini client not initialized. Please set GEMINI_API_KEY environment variable."
+
+    return get_explanation(
+        state["image"],
+        state["bboxes"],
+        state["score"],
+        state["category"],
+        gemini_client
+    )
+
+
+def create_ui(gemini_client):
+    """
+    Create and configure the Gradio UI.
+
+    Args:
+        gemini_client: Initialized Gemini client
+
+    Returns:
+        Gradio Blocks interface
+    """
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🔍 Industrial Anomaly Detection")
+        gr.Markdown("PatchCore + EfficientAD + Explainable AI")
+
+        state = gr.State()
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(label="Upload Image", type="numpy", height=300)
+
+                model_dropdown = gr.Dropdown(
+                    choices=["patchcore", "efficientad"],
+                    value="patchcore",
+                    label="Model"
+                )
+
+                category_dropdown = gr.Dropdown(
+                    choices=list(THRESHOLDS["patchcore"].keys()),
+                    value="bottle",
+                    label="Category"
+                )
+
+                detect_btn = gr.Button("🚀 Run Detection")
+                explain_btn = gr.Button("🧠 Explain Anomaly")
+
+            with gr.Column(scale=2):
+                with gr.Row():
+                    out_original = gr.Image(label="Original")
+                    out_heatmap = gr.Image(label="Heatmap")
+
+                with gr.Row():
+                    out_overlay = gr.Image(label="Overlay")
+                    out_mask = gr.Image(label="Predicted Mask")
+
+                with gr.Row():
+                    score_box = gr.Textbox(label="Score")
+                    threshold_box = gr.Textbox(label="Threshold")
+                    status_box = gr.Textbox(label="Status")
+
+                explanation_box = gr.Textbox(label="Explanation", lines=3)
+
+        # Button actions
+        detect_btn.click(
+            fn=lambda img, model, cat: detect(img, model, cat, gemini_client),
+            inputs=[input_image, model_dropdown, category_dropdown],
+            outputs=[
+                out_original,
+                out_heatmap,
+                out_overlay,
+                out_mask,
+                score_box,
+                threshold_box,
+                status_box,
+                state
+            ],
+        )
+
+        explain_btn.click(
+            fn=explain,
+            inputs=[state],
+            outputs=explanation_box
+        )
+
+    return demo
+
+
+def main():
+    """Main entry point for the application."""
+    import os
+    from dotenv import load_dotenv
+
+    # Load environment variables
+    load_dotenv()
+
+    # Get Gemini API key
+    api_key = os.getenv("GEMINI_API_KEY")
+    if not api_key:
+        raise ValueError(
+            "GEMINI_API_KEY not found. Please set it in .env file or environment variables."
+        )
+
+    # Initialize Gemini client
+    gemini_client = init_gemini_client(api_key)
+
+    # Create and launch UI
+    demo = create_ui(gemini_client)
+    
+    # Configure Gradio settings from environment variables
+    share = os.getenv("GRADIO_SHARE", "False").lower() == "true"
+    debug = os.getenv("GRADIO_DEBUG", "False").lower() == "true"
+    
+    demo.launch(share=share, debug=debug)
+
+
+if __name__ == "__main__":
+    main()

config.py

@@ -0,0 +1,85 @@
+"""
+Configuration settings for the anomaly detection project.
+
+This module contains all configurable parameters including thresholds,
+model mappings, and other constants used throughout the application.
+"""
+
+# Model directory mapping
+MODEL_TO_DIR = {
+    "patchcore": "Patchcore",
+    "efficientad": "EfficientAd",
+}
+
+# MVTec AD dataset categories
+MVTEC_CATEGORIES = [
+    "bottle", "cable", "capsule", "carpet", "grid",
+    "hazelnut", "leather", "metal_nut", "pill", "screw",
+    "tile", "toothbrush", "transistor", "wood", "zipper"
+]
+
+# Precomputed thresholds for each model and category
+# These thresholds are computed at the 95th percentile of normal training scores
+THRESHOLDS = {
+    "patchcore": {
+        "bottle": 0.3218444108963013,
+        "cable": 0.34408192038536073,
+        "capsule": 0.5454285681247711,
+        "carpet": 0.3088440954685211,
+        "grid": 0.25913039445877073,
+        "hazelnut": 0.10068576037883759,
+        "leather": 0.2726534068584442,
+        "metal_nut": 0.34413049668073653,
+        "pill": 0.26968240439891816,
+        "screw": 0.49187072515487673,
+        "tile": 0.3581161931157112,
+        "toothbrush": 0.3721309259533882,
+        "transistor": 0.45495494604110714,
+        "wood": 0.1711873710155487,
+        "zipper": 0.4981046631932258
+    },
+    "efficientad": {
+        "bottle": 0.49928921461105347,
+        "cable": 0.4673861160874367,
+        "capsule": 0.5370000839233399,
+        "carpet": 0.49847708493471143,
+        "grid": 0.5295769184827804,
+        "hazelnut": 0.5202932059764862,
+        "leather": 0.504090940952301,
+        "metal_nut": 0.5047085165977478,
+        "pill": 0.5043391764163971,
+        "screw": 0.7167768508195878,
+        "tile": 0.5030474990606308,
+        "toothbrush": 0.5439804702997207,
+        "transistor": 0.5076832294464111,
+        "wood": 0.5024313390254974,
+        "zipper": 1.0
+    }
+}
+
+# Hugging Face repository ID for checkpoints
+HF_REPO_ID = "micguida1/mvtec-anomaly-checkpoints"
+
+# Image size for model input
+IMAGE_SIZE = 256
+
+# Threshold multiplier for sensitivity adjustment
+# Lower value = more sensitive (more anomalies detected)
+THRESHOLD_MULTIPLIER = 0.85
+
+# Visualization parameters
+HEATMAP_ALPHA = 0.5
+OVERLAY_ALPHA = 0.5
+
+# PatchCore specific parameters
+PATCHCORE_BINARY_THRESHOLD = 0.60
+PATCHCORE_MIN_CONTOUR_AREA = 100
+PATCHCORE_MAX_INTENSITY_THRESHOLD = 0.75
+PATCHCORE_BLUR_KERNEL = (7, 7)
+PATCHCORE_MORPH_KERNEL = (5, 5)
+PATCHCORE_FG_THRESHOLD = 15
+PATCHCORE_FG_MORPH_KERNEL = (9, 9)
+
+# EfficientAD specific parameters
+EFFICIENTAD_BINARY_THRESHOLD = 0.5
+EFFICIENTAD_MIN_CONTOUR_AREA = 5

core/explain.py

@@ -0,0 +1,97 @@
+"""
+Explainability module using Gemini VLM.
+
+This module provides functions to generate human-readable explanations
+for detected anomalies using Google's Gemini Vision Language Model.
+"""
+
+from PIL import Image
+import cv2
+from google import genai
+import numpy as np
+
+# Model configuration
+GEMINI_MODEL = "gemini-flash-lite-latest"
+
+
+def get_explanation(
+    original_image: np.ndarray,
+    bboxes: list,
+    score: float,
+    category: str,
+    client
+) -> str:
+    """
+    Generate an explanation for the detected anomaly using Gemini VLM.
+
+    Args:
+        original_image: Original input image in RGB format
+        bboxes: List of bounding boxes [x1, y1, x2, y2] in 256x256 scale
+        score: Anomaly score
+        category: MVTec AD category
+        client: Initialized Gemini API client
+
+    Returns:
+        Explanation text from the model
+    """
+    if not bboxes:
+        return "No anomaly detected."
+
+    # Scale bounding boxes from 256x256 to original image size
+    h_orig, w_orig = original_image.shape[:2]
+    scale_x = w_orig / 256.0
+    scale_y = h_orig / 256.0
+
+    # Draw red bounding boxes on a copy of the original image
+    annotated_img = original_image.copy()
+    for (x1, y1, x2, y2) in bboxes:
+        x1_s, y1_s = int(x1 * scale_x), int(y1 * scale_y)
+        x2_s, y2_s = int(x2 * scale_x), int(y2 * scale_y)
+
+        # Dynamic thickness based on image size
+        thickness = max(2, int(max(h_orig, w_orig) * 0.005))
+        cv2.rectangle(annotated_img, (x1_s, y1_s), (x2_s, y2_s), (255, 0, 0), thickness)
+
+    # Convert to PIL Image
+    annotated_pil = Image.fromarray(annotated_img)
+
+    # Construct prompt
+    prompt = f"""
+    You are an expert industrial quality control inspector.
+    We are inspecting a: {category}
+
+    An anomaly detection model has flagged a potential defect, highlighted by the RED BOUNDING BOX in the provided image.
+
+    Your task is to classify the defect inside the red box and assess its severity.
+    Common defects for {category} include: scratches, cuts, cracks, holes, structural damage, or severe discoloration.
+
+    Analyze the highlighted region carefully in the context of the whole object.
+
+    Only Provide your final assessment strictly in this format:
+    Defect: <Short name, e.g., Deep Scratch, Surface Cut, Crack, Contamination, Colouration>
+    Location: <Where is it on the object?>
+    Severity: <Low/Medium/High>
+    """
+
+    # Generate response from Gemini with error handling
+    try:
+        response = client.models.generate_content(
+            model=GEMINI_MODEL,
+            contents=[prompt, annotated_pil]
+        )
+        return response.text
+    except Exception as e:
+        return f"Failed to generate explanation: {str(e)}"
+
+
+def init_gemini_client(api_key: str):
+    """
+    Initialize the Gemini API client.
+
+    Args:
+        api_key: Gemini API key
+
+    Returns:
+        Initialized genai client
+    """
+    return genai.Client(api_key=api_key)

core/inference.py

@@ -0,0 +1,98 @@
+"""
+Inference module for anomaly detection.
+
+This module handles image preprocessing, model inference, and basic output extraction.
+It does not include postprocessing logic for mask generation.
+"""
+
+import cv2
+import numpy as np
+import torch
+
+from config import IMAGE_SIZE
+
+
+def preprocess_image(image: np.ndarray) -> torch.Tensor:
+    """
+    Preprocess an image for model input.
+
+    Args:
+        image: Input image in RGB format (H, W, 3)
+
+    Returns:
+        Preprocessed tensor ready for model inference (1, 3, H, W)
+    """
+    # Resize to model input size
+    img_resized = cv2.resize(image, (IMAGE_SIZE, IMAGE_SIZE))
+
+    # Normalize to [0, 1]
+    img = img_resized / 255.0
+
+    # Transpose from (H, W, C) to (C, H, W)
+    img = np.transpose(img, (2, 0, 1))
+
+    # Add batch dimension
+    img = np.expand_dims(img, axis=0)
+
+    # Convert to tensor
+    tensor = torch.tensor(img, dtype=torch.float32)
+
+    return tensor
+
+
+def run_inference(model, image: np.ndarray, model_name: str, category: str):
+    """
+    Run inference on a single image.
+
+    Args:
+        model: Loaded anomaly detection model
+        image: Input image in RGB format
+        model_name: Name of the model being used
+        category: MVTec AD category
+
+    Returns:
+        tuple: (heatmap, pred_mask_raw, score)
+            - heatmap: Raw anomaly heatmap (H, W)
+            - pred_mask_raw: Raw predicted mask if available (H, W) or None
+            - score: Anomaly score (float)
+    """
+    # Preprocess the image
+    tensor = preprocess_image(image)
+
+    # Get device from model
+    device = next(model.parameters()).device
+    tensor = tensor.to(device)
+
+    # Run inference
+    with torch.no_grad():
+        output = model(tensor)
+
+    # Extract outputs based on output format with validation
+    pred_mask_raw = None
+
+    if hasattr(output, "anomaly_map") and hasattr(output, "pred_score"):
+        heatmap = output.anomaly_map
+        score = output.pred_score
+        pred_mask_raw = getattr(output, "pred_mask", None)
+    elif isinstance(output, dict) and "anomaly_map" in output and "pred_score" in output:
+        heatmap = output["anomaly_map"]
+        score = output["pred_score"]
+        pred_mask_raw = output.get("pred_mask", None)
+    elif isinstance(output, tuple) and len(output) >= 2:
+        score, heatmap = output[0], output[1]
+    else:
+        raise ValueError(
+            f"Model output must contain anomaly_map and pred_score. "
+            f"Got output type: {type(output)}. "
+            f"If using a dict, ensure it has 'anomaly_map' and 'pred_score' keys. "
+            f"If using an object, ensure it has 'anomaly_map' and 'pred_score' attributes."
+        )
+
+    # Convert to numpy
+    heatmap = heatmap.squeeze().cpu().numpy()
+    score = float(score.cpu().numpy() if torch.is_tensor(score) else score)
+
+    if pred_mask_raw is not None:
+        pred_mask_raw = pred_mask_raw.squeeze().cpu().numpy()
+
+    return heatmap, pred_mask_raw, score

core/postprocess.py

@@ -0,0 +1,180 @@
+"""
+Postprocessing module for anomaly detection.
+
+This module handles mask generation, bounding box extraction, and validation
+from raw anomaly heatmaps. It contains model-specific logic for both PatchCore
+and EfficientAD.
+"""
+
+import cv2
+import numpy as np
+
+from config import (
+    PATCHCORE_BINARY_THRESHOLD,
+    PATCHCORE_MIN_CONTOUR_AREA,
+    PATCHCORE_MAX_INTENSITY_THRESHOLD,
+    PATCHCORE_BLUR_KERNEL,
+    PATCHCORE_MORPH_KERNEL,
+    PATCHCORE_FG_THRESHOLD,
+    PATCHCORE_FG_MORPH_KERNEL,
+    EFFICIENTAD_BINARY_THRESHOLD,
+    EFFICIENTAD_MIN_CONTOUR_AREA,
+)
+
+
+def postprocess_patchcore(heatmap: np.ndarray, original_image: np.ndarray, is_anomaly: bool):
+    """
+    Postprocess heatmap for PatchCore model.
+
+    Args:
+        heatmap: Raw anomaly heatmap (H, W)
+        original_image: Original resized image (H, W, 3)
+        is_anomaly: Whether the image is classified as an anomaly
+
+    Returns:
+        tuple: (final_mask, bboxes, heatmap_vis)
+            - final_mask: Binary mask of anomaly regions (uint8)
+            - bboxes: List of bounding boxes [x1, y1, x2, y2]
+            - heatmap_vis: Normalized heatmap for visualization (H, W)
+    """
+    h, w = heatmap.shape
+
+    # Blur the heatmap for smoother contours
+    heatmap_blurred = cv2.GaussianBlur(heatmap, PATCHCORE_BLUR_KERNEL, 0)
+
+    # Normalize to [0, 1] with robust handling of constant heatmaps
+    h_min = float(heatmap_blurred.min())
+    h_max = float(heatmap_blurred.max())
+    h_range = h_max - h_min
+    if h_range < 1e-6:
+        # Heatmap is essentially constant, normalize to all zeros
+        heatmap_vis = np.zeros_like(heatmap_blurred)
+    else:
+        heatmap_vis = (heatmap_blurred - h_min) / h_range
+
+    # Foreground masking to ignore background
+    gray = cv2.cvtColor(original_image, cv2.COLOR_RGB2GRAY)
+    _, fg_mask = cv2.threshold(gray, PATCHCORE_FG_THRESHOLD, 255, cv2.THRESH_BINARY)
+    kernel_fg = np.ones(PATCHCORE_FG_MORPH_KERNEL, np.uint8)
+    fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, kernel_fg)
+
+    heatmap_vis[fg_mask == 0] = 0
+
+    # Initialize outputs
+    final_mask = np.zeros_like(heatmap_vis, dtype=np.uint8)
+    bboxes = []
+
+    if is_anomaly:
+        # Threshold to binary
+        binary = (heatmap_vis > PATCHCORE_BINARY_THRESHOLD).astype(np.uint8) * 255
+
+        # Morphological operations to clean up
+        kernel = np.ones(PATCHCORE_MORPH_KERNEL, np.uint8)
+        binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
+        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
+
+        # Find contours
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        # Filter valid contours based on area and intensity
+        valid_contours = []
+        for c in contours:
+            x, y, cw, ch = cv2.boundingRect(c)
+            area = cw * ch
+
+            if area < PATCHCORE_MIN_CONTOUR_AREA:
+                continue
+
+            # Check max intensity within contour
+            mask_temp = np.zeros_like(heatmap_vis, dtype=np.uint8)
+            cv2.drawContours(mask_temp, [c], -1, 1, thickness=-1)
+            max_intensity = heatmap_vis[mask_temp == 1].max()
+
+            if max_intensity > PATCHCORE_MAX_INTENSITY_THRESHOLD:
+                valid_contours.append(c)
+
+        # Draw valid contours and extract bounding boxes
+        if valid_contours:
+            valid_contours = sorted(valid_contours, key=cv2.contourArea, reverse=True)
+            cv2.drawContours(final_mask, valid_contours, -1, 255, thickness=-1)
+
+            for c in valid_contours:
+                x, y, cw, ch = cv2.boundingRect(c)
+                bboxes.append([x, y, x + cw, y + ch])
+
+    return final_mask, bboxes, heatmap_vis
+
+
+def postprocess_efficientad(heatmap: np.ndarray, is_anomaly: bool):
+    """
+    Postprocess heatmap for EfficientAD model.
+
+    Args:
+        heatmap: Raw anomaly heatmap (H, W)
+        is_anomaly: Whether the image is classified as an anomaly
+
+    Returns:
+        tuple: (final_mask, bboxes, heatmap_vis)
+            - final_mask: Binary mask of anomaly regions (uint8)
+            - bboxes: List of bounding boxes [x1, y1, x2, y2]
+            - heatmap_vis: Normalized heatmap for visualization (H, W)
+    """
+    h, w = heatmap.shape
+
+    # Normalize with adaptive strategy and robust handling
+    amap_min = float(heatmap.min())
+    amap_max = float(heatmap.max())
+    amap_range = amap_max - amap_min
+
+    if amap_range < 0.1:
+        if amap_range > 1e-6:
+            heatmap_vis = (heatmap - amap_min) / amap_range
+        else:
+            # Heatmap is essentially constant
+            heatmap_vis = np.zeros_like(heatmap)
+    else:
+        # Clip to [0, 1] range, assuming heatmap is already roughly normalized
+        heatmap_vis = np.clip(heatmap, 0, 1)
+
+    # Dim the heatmap if not an anomaly
+    if not is_anomaly:
+        heatmap_vis = heatmap_vis * 0.3
+
+    # Initialize outputs
+    final_mask = np.zeros_like(heatmap, dtype=np.uint8)
+    bboxes = []
+
+    if is_anomaly:
+        # Threshold exactly at 0.5
+        binary = (heatmap_vis > EFFICIENTAD_BINARY_THRESHOLD).astype(np.uint8) * 255
+
+        # Find contours
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        # Filter by minimum area
+        for c in contours:
+            x, y, cw, ch = cv2.boundingRect(c)
+            if cw * ch > EFFICIENTAD_MIN_CONTOUR_AREA:
+                cv2.drawContours(final_mask, [c], -1, 255, thickness=-1)
+                bboxes.append([x, y, x + cw, y + ch])
+
+    return final_mask, bboxes, heatmap_vis
+
+
+def postprocess(heatmap: np.ndarray, original_image: np.ndarray, model_name: str, is_anomaly: bool):
+    """
+    Main postprocessing function that routes to the appropriate model-specific processor.
+
+    Args:
+        heatmap: Raw anomaly heatmap (H, W)
+        original_image: Original resized image (H, W, 3)
+        model_name: Name of the model ("patchcore" or "efficientad")
+        is_anomaly: Whether the image is classified as an anomaly
+
+    Returns:
+        tuple: (final_mask, bboxes, heatmap_vis)
+    """
+    if model_name == "efficientad":
+        return postprocess_efficientad(heatmap, is_anomaly)
+    else:
+        return postprocess_patchcore(heatmap, original_image, is_anomaly)

core/visualization.py

@@ -0,0 +1,146 @@
+"""
+Visualization module for anomaly detection results.
+
+This module provides functions to create visual outputs including heatmaps,
+overlays, and predicted mask visualizations.
+"""
+
+import cv2
+import numpy as np
+
+from config import HEATMAP_ALPHA, OVERLAY_ALPHA
+
+
+def create_overlay(image: np.ndarray, heatmap: np.ndarray, model_name: str) -> np.ndarray:
+    """
+    Create an overlay of the heatmap on the original image.
+
+    Args:
+        image: Original image in RGB format (H, W, 3)
+        heatmap: Normalized heatmap (H, W) in range [0, 1]
+        model_name: Name of the model for model-specific handling
+
+    Returns:
+        Overlay image in RGB format (H, W, 3)
+    """
+    image_resized = cv2.resize(image, (256, 256))
+
+    # Convert heatmap to uint8 and apply colormap
+    heatmap_uint8 = (heatmap * 255).astype(np.uint8)
+    heatmap_color = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
+    heatmap_color = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB)
+
+    # Model-specific handling for EfficientAD padding
+    if model_name == "efficientad":
+        # Mask out zero values (padding) to show original image
+        mask_0 = (heatmap == 0)[..., np.newaxis]
+        overlay = np.where(mask_0, image_resized, cv2.addWeighted(image_resized, OVERLAY_ALPHA, heatmap_color, HEATMAP_ALPHA, 0))
+    else:
+        overlay = cv2.addWeighted(image_resized, OVERLAY_ALPHA, heatmap_color, HEATMAP_ALPHA, 0)
+
+    return overlay
+
+
+def create_mask_visualization(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """
+    Create a visualization of the predicted mask overlaid on the image.
+
+    Args:
+        image: Original image in RGB format (H, W, 3)
+        mask: Binary mask (H, W) where non-zero values indicate anomaly
+
+    Returns:
+        Visualization image with semi-transparent red mask and contours
+    """
+    image_resized = cv2.resize(image, (256, 256))
+    vis_img = image_resized.copy()
+
+    if np.any(mask):
+        # Create a red color mask
+        color_mask = np.zeros_like(image_resized)
+        color_mask[mask > 0] = [255, 0, 0]  # RGB Red
+
+        # Blend with original image
+        vis_img = cv2.addWeighted(vis_img, 0.7, color_mask, 0.3, 0)
+
+        # Draw contours
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(vis_img, contours, -1, (255, 255, 255), 2)
+
+    return vis_img
+
+
+def draw_bounding_boxes(overlay: np.ndarray, mask_vis: np.ndarray, bboxes: list):
+    """
+    Draw bounding boxes on both overlay and mask visualization images.
+
+    Args:
+        overlay: Overlay image to draw on (modified in-place)
+        mask_vis: Mask visualization image to draw on (modified in-place)
+        bboxes: List of bounding boxes [x1, y1, x2, y2]
+    """
+    for (x1, y1, x2, y2) in bboxes:
+        # Green boxes on overlay
+        cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        # Blue boxes on mask visualization
+        cv2.rectangle(mask_vis, (x1, y1), (x2, y2), (255, 0, 0), 2)
+
+
+def create_heatmap_color(heatmap: np.ndarray, model_name: str) -> np.ndarray:
+    """
+    Create a colored heatmap image suitable for display.
+
+    Args:
+        heatmap: Normalized heatmap (H, W) in range [0, 1]
+        model_name: Name of the model for model-specific handling
+
+    Returns:
+        Colored heatmap in RGB format (H, W, 3)
+    """
+    heatmap_uint8 = (heatmap * 255).astype(np.uint8)
+    heatmap_color = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
+    heatmap_color = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB)
+
+    # For EfficientAD, make padding (zero values) black
+    if model_name == "efficientad":
+        heatmap_color[heatmap == 0] = [0, 0, 0]
+
+    return heatmap_color
+
+
+def create_visuals(
+    image: np.ndarray,
+    heatmap: np.ndarray,
+    mask: np.ndarray,
+    bboxes: list,
+    model_name: str
+) -> tuple:
+    """
+    Create all visualization outputs for a single inference result.
+
+    Args:
+        image: Original input image in RGB format
+        heatmap: Normalized heatmap (H, W)
+        mask: Binary mask (H, W)
+        bboxes: List of bounding boxes
+        model_name: Name of the model
+
+    Returns:
+        tuple: (original_resized, heatmap_color, overlay, mask_vis)
+    """
+    # Resize original image to 256x256
+    original_resized = cv2.resize(image, (256, 256))
+
+    # Create heatmap visualization
+    heatmap_color = create_heatmap_color(heatmap, model_name)
+
+    # Create overlay
+    overlay = create_overlay(image, heatmap, model_name)
+
+    # Create mask visualization
+    mask_vis = create_mask_visualization(image, mask)
+
+    # Draw bounding boxes on overlay and mask visualization
+    draw_bounding_boxes(overlay, mask_vis, bboxes)
+
+    return original_resized, heatmap_color, overlay, mask_vis

models/model_loader.py

@@ -0,0 +1,95 @@
+"""
+Model loading and caching module.
+
+This module provides functions to load anomaly detection models from
+Hugging Face Hub with caching support to avoid reloading the same model multiple times.
+"""
+
+import os
+import torch
+from collections import OrderedDict
+from huggingface_hub import hf_hub_download
+from anomalib.models import Patchcore, EfficientAd
+
+from config import HF_REPO_ID, MODEL_TO_DIR
+
+# Maximum number of models to keep in cache (prevents unbounded memory growth)
+MAX_MODEL_CACHE_SIZE = 30
+
+# Global model cache with LRU eviction (using OrderedDict)
+_model_cache = OrderedDict()
+
+
+def get_ckpt_path(model_name: str, category: str) -> str:
+    """
+    Download or retrieve the checkpoint file for a given model and category.
+
+    Args:
+        model_name: Name of the model ("patchcore" or "efficientad")
+        category: MVTec AD category (e.g., "bottle", "cable")
+
+    Returns:
+        Path to the downloaded checkpoint file
+    """
+    dirname = MODEL_TO_DIR[model_name]
+    hf_path = f"{dirname}/MVTecAD/{category}/latest/weights/lightning/model.ckpt"
+
+    return hf_hub_download(
+        repo_id=HF_REPO_ID,
+        filename=hf_path,
+        local_dir="models",
+        local_dir_use_symlinks=False,
+    )
+
+
+def load_model(model_name: str, category: str):
+    """
+    Load an anomaly detection model with caching and LRU eviction.
+
+    Args:
+        model_name: Name of the model ("patchcore" or "efficientad")
+        category: MVTec AD category
+
+    Returns:
+        Loaded model on the appropriate device (CUDA if available)
+
+    Raises:
+        ValueError: If an unknown model name is provided
+    """
+    key = f"{model_name}_{category}"
+
+    # Return cached model if available (move to end to mark as recently used)
+    if key in _model_cache:
+        _model_cache.move_to_end(key)
+        return _model_cache[key]
+
+    # Evict least recently used model if cache is full
+    if len(_model_cache) >= MAX_MODEL_CACHE_SIZE:
+        _model_cache.popitem(last=False)  # Remove first (oldest) item
+
+    # Download checkpoint
+    ckpt = get_ckpt_path(model_name, category)
+
+    # Load the appropriate model type
+    if model_name == "patchcore":
+        model = Patchcore.load_from_checkpoint(ckpt)
+    elif model_name == "efficientad":
+        model = EfficientAd.load_from_checkpoint(ckpt)
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+
+    # Set evaluation mode and move to device
+    model.eval()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+
+    # Cache the model (add to end)
+    _model_cache[key] = model
+
+    return model
+
+
+def clear_model_cache():
+    """Clear the model cache to free memory."""
+    global _model_cache
+    _model_cache.clear()

requirements.txt

@@ -0,0 +1,37 @@
+# Anomaly Detection Dependencies
+
+# Core ML Framework
+torch>=2.0.0
+torchvision
+
+# Anomalib for pre-trained models
+anomalib==2.2.0
+
+# Gradio for UI
+gradio>=4.0.0
+
+# OpenCV for image processing
+opencv-python-headless>=4.8.0
+
+# Image processing
+Pillow>=10.0.0
+matplotlib>=3.7.0
+
+# Hugging Face Hub for model downloads
+huggingface-hub>=0.19.0
+huggingface-hub[cli]
+
+# Google Gemini for explainable AI
+google-genai>=0.3.0
+
+# Utilities
+numpy>=1.24.0
+tqdm>=4.65.0
+python-dotenv>=1.0.0
+
+# Training support
+tensorboard>=2.16.0
+
+# Optional: ONNX support (if needed)
+# onnx>=1.14.0
+# openvino>=2023.3.0