app.py

import tempfile
import time
from functools import lru_cache
from pathlib import Path

import gradio as gr
import torch
from anomalib.data import MVTecAD
from anomalib.engine import Engine
from anomalib.metrics import AUROC, Evaluator
from anomalib.models import Patchcore
from anomalib.visualization import visualize_anomaly_map, visualize_mask
from anomalib.visualization.image.functional import overlay_images
from safetensors import safe_open
from safetensors.torch import save_file
from torchvision.transforms import v2

ROOT = Path(__file__).parent
DATASET = ROOT / "dataset"
CACHE = ROOT / ".cache"
RATIO = 0.01
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

CATEGORIES = ['bottle', 'cable', 'capsule', 'carpet', 'grid', 'hazelnut', 'leather', 'metal_nut',
              'pill', 'screw', 'tile', 'toothbrush', 'transistor', 'wood', 'zipper']

PAPER = dict(zip(CATEGORIES, zip(
    [100, 99.3, 98.0, 98.0, 98.6, 100, 100, 99.7, 97.0, 96.4, 99.4, 100, 99.9, 99.2, 99.2],
    [98.5, 98.2, 98.8, 98.9, 98.6, 98.6, 99.3, 98.4, 97.1, 99.2, 96.1, 98.5, 94.9, 95.1, 98.8])))

TO_TENSOR = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])


def build(category):
    evaluator = Evaluator(test_metrics=[
        AUROC(fields=['pred_score', 'gt_label'], prefix='image_'),
        AUROC(fields=['anomaly_map', 'gt_mask'], prefix='pixel_'),
    ])
    module = Patchcore(coreset_sampling_ratio=RATIO, evaluator=evaluator)
    # python 3.14 defaults to forkserver, which re-imports __main__ in every worker
    data = MVTecAD(root=DATASET, category=category, num_workers=0)

    with tempfile.TemporaryDirectory() as tmp:
        engine = Engine(default_root_dir=tmp, logger=False)
        engine.fit(model=module, datamodule=data)
        metrics = engine.test(model=module, datamodule=data)[0]

    tensors = {'memory_bank': module.model.memory_bank.half().cpu()}
    tensors |= {k: v.cpu() for k, v in module.post_processor.state_dict().items()}
    CACHE.mkdir(exist_ok=True)
    save_file(tensors, CACHE / f"{category}.safetensors",
              metadata={k: f"{v:.4f}" for k, v in metrics.items()})


@lru_cache
def load(category):
    path = CACHE / f"{category}.safetensors"
    if not path.exists():
        build(category)

    module = Patchcore(coreset_sampling_ratio=RATIO)
    with safe_open(path, framework='pt') as f:
        metrics = f.metadata()
        state = {k: f.get_tensor(k) for k in f.keys()}  # noqa: SIM118 (not a dict)

    module.model.memory_bank = state.pop('memory_bank').float()
    module.post_processor.load_state_dict(state)
    return module.eval().to(DEVICE), metrics


@torch.no_grad()
def infer(module, image):
    tensor = module.pre_processor.transform(TO_TENSOR(image))[None].to(DEVICE)
    start = time.perf_counter()
    out = module.post_processor(module.model(tensor))
    elapsed = (time.perf_counter() - start) * 1000
    return float(out.pred_score), out.anomaly_map[0].cpu(), out.pred_mask[0].cpu(), elapsed


def samples(category):
    test = DATASET / category / 'test'
    good = sorted((test / 'good').glob('*.png'))[:1]
    defects = [min(d.glob('*.png')) for d in sorted(test.iterdir()) if d.name != 'good']
    return [(str(p), p.parent.name) for p in good + defects[:3]]


def scoreboard(category, metrics, threshold):
    auroc, pw_auroc = PAPER[category]
    return (f"| test set | ours | paper |\n|---|---|---|\n"
            f"| image AUROC | {float(metrics['image_AUROC']) * 100:.1f} | {auroc:.1f} |\n"
            f"| pixel AUROC | {float(metrics['pixel_AUROC']) * 100:.1f} | {pw_auroc:.1f} |\n\n"
            f"F1-optimal threshold on the validation split: {float(threshold):.1f}, which the "
            f"min-max normalization maps to 0.5.")


def select(category, progress=gr.Progress()):  # noqa: B008
    progress(0, desc=f"Fitting {category}" if not (CACHE / f"{category}.safetensors").exists()
             else f"Loading {category}")
    module, metrics = load(category)
    return (gr.update(value=samples(category)),
            scoreboard(category, metrics, module.post_processor.image_threshold))


def pick(category, evt: gr.SelectData):
    return samples(category)[evt.index][0]


def verdict(score, threshold):
    return f"FAIL - score {score:.3f}" if score >= threshold else f"PASS - score {score:.3f}"


def run(category, image, threshold):
    if image is None:
        return None, None, '', '', 0.0
    module, _ = load(category)
    score, anomaly_map, mask, elapsed = infer(module, image)
    # upscaled to the input resolution, otherwise gradio renders them at their native 256px
    size = image.size[::-1]
    heatmap = visualize_anomaly_map(v2.functional.resize(anomaly_map, size))
    mask = v2.functional.resize(mask, size, interpolation=v2.InterpolationMode.NEAREST)
    contour = visualize_mask(mask, mode='contour', color=(255, 0, 0))
    return (heatmap, overlay_images(image, contour),
            verdict(score, threshold), f"{elapsed:.0f} ms on {DEVICE}", score)


with gr.Blocks(title="PatchCore Inspector") as demo:
    gr.Markdown("## PatchCore Inspector\nAnomaly detection on MVTec AD, "
                "memory bank subsampled to 1% with a greedy coreset.")
    score = gr.State(0.0)

    with gr.Row():
        category = gr.Dropdown(CATEGORIES, value=CATEGORIES[0], label="Product")
        result = gr.Textbox(label="Verdict", interactive=False)
        latency = gr.Textbox(label="Inference time", interactive=False)
        threshold = gr.Slider(0, 1, 0.5, step=0.01, label="Decision threshold", scale=2)

    with gr.Row(equal_height=True):
        image = gr.Image(type='pil', sources=['upload'], label="Image under inspection", height=500)
        heatmap = gr.Image(label="Anomaly map", height=500)
        overlay = gr.Image(label="Overlay", height=500)

    with gr.Row():
        gallery = gr.Gallery(samples(CATEGORIES[0]), label="Test samples", columns=4, height=290,
                             object_fit='contain', allow_preview=False, scale=2)
        scores = gr.Markdown()

    outputs = [heatmap, overlay, result, latency, score]
    category.change(select, category, [gallery, scores]).then(run, [category, image, threshold], outputs)
    image.change(run, [category, image, threshold], outputs)
    gallery.select(pick, category, image)
    threshold.change(verdict, [score, threshold], result)
    demo.load(select, category, [gallery, scores])

if __name__ == '__main__':
    demo.launch()