inference.py

import sys
import warnings
import os
import signal
import atexit
import time

os.environ["PYTHONWARNINGS"] = "ignore"
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
warnings.filterwarnings("ignore")


def _force_exit(signum=None, frame=None):
    os._exit(1)


signal.signal(signal.SIGINT, _force_exit)
signal.signal(signal.SIGTERM, _force_exit)
atexit.register(lambda: os._exit(0) if hasattr(os, '_exit') else None)

import torch
import torch.nn as nn
import numpy as np
import random
from functools import partial

torch._inductor.config.combo_kernels = False

_dinomaly_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Dinomaly')
sys.path.insert(0, _dinomaly_dir)
os.chdir(_dinomaly_dir)

from models.uad import ViTill
from models.vision_transformer import Block as VitBlock, bMlp, LinearAttention2
from dataset import MVTecDataset, get_data_transforms
from utils import evaluation_batch
from dinov3.hub.backbones import load_dinov3_model


DATA_PATH = 'c:/Users/IDEXX_BMB/Desktop/New folder/Data/mvtec_anomaly_detection'
WEIGHTS_PATH = os.path.join(_dinomaly_dir, 'saved_results', 'bottle_model.pth')
DEVICE = 'cuda:0'
BATCH_SIZE = 16
IMAGE_SIZE = 512       # DINOv3 patch16: match cnulab config
CROP_SIZE = 448        # 448 / 16 = 28x28 = 784 patches
WARMUP_ITERS = 10
TIMED_ITERS = 100


def try_compile(model):
    """torch.compile with inductor reduce-overhead (CUDAGraphs enabled via RoPE caching)."""
    try:
        return torch.compile(model, mode='reduce-overhead')
    except Exception:
        try:
            return torch.compile(model, mode='default')
        except Exception:
            return model


def build_model():
    """Build Dinomaly model with DINOv3 encoder."""
    encoder_name = 'dinov3_vitb16'
    encoder_weight = os.path.join(_dinomaly_dir, 'weights', 'dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth')
    target_layers = [2, 3, 4, 5, 6, 7, 8, 9]
    encoder = load_dinov3_model(encoder_name, layers_to_extract_from=target_layers, pretrained_weight_path=encoder_weight)

    embed_dim, num_heads = 768, 12
    fuse_layer_encoder = [[0, 1, 2, 3], [4, 5, 6, 7]]
    fuse_layer_decoder = [[0, 1, 2, 3], [4, 5, 6, 7]]

    bottleneck = nn.ModuleList([bMlp(embed_dim, embed_dim * 4, embed_dim, drop=0.2)])
    decoder = nn.ModuleList([
        VitBlock(dim=embed_dim, num_heads=num_heads, mlp_ratio=4.,
                 qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-8),
                 attn_drop=0., attn=LinearAttention2)
        for _ in range(8)
    ])

    model = ViTill(encoder=encoder, bottleneck=bottleneck, decoder=decoder,
                   target_layers=target_layers, mask_neighbor_size=0,
                   fuse_layer_encoder=fuse_layer_encoder,
                   fuse_layer_decoder=fuse_layer_decoder, recentering=False)
    return model


def benchmark_speed(model, dummy_input):
    """Warmup + timed forward passes under AMP autocast. Returns avg ms per image."""
    model.eval()
    with torch.inference_mode():
        for _ in range(WARMUP_ITERS):
            with torch.amp.autocast('cuda', dtype=torch.bfloat16):
                model(dummy_input)
        torch.cuda.synchronize()

        start = time.perf_counter()
        for _ in range(TIMED_ITERS):
            with torch.amp.autocast('cuda', dtype=torch.bfloat16):
                model(dummy_input)
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - start

    ms_per_image = (elapsed / TIMED_ITERS / BATCH_SIZE) * 1000
    print(f'  {TIMED_ITERS} iters x {BATCH_SIZE} imgs: {elapsed:.2f}s total, {ms_per_image:.2f} ms/image')
    return ms_per_image


class AutocastWrapper(nn.Module):
    """Wraps a model so evaluation_batch runs under AMP autocast."""
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, x):
        with torch.inference_mode(), torch.amp.autocast('cuda', dtype=torch.bfloat16):
            en, de = self.model(x)
        return [e.float() for e in en], [d.float() for d in de]


if __name__ == '__main__':
    torch.manual_seed(1)
    np.random.seed(1)
    random.seed(1)
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True
    torch.set_float32_matmul_precision('high')

    print('=' * 60)
    print('Inference Benchmark — Dinomaly DINOv3 (bottle, AMP-BF16)')
    print('=' * 60)

    if not os.path.exists(WEIGHTS_PATH):
        print(f'\nERROR: Trained weights not found at {WEIGHTS_PATH}')
        print('Run train_dinov3.py first to train and save the model.')
        sys.exit(1)

    # Load test data
    data_transform, gt_transform = get_data_transforms(IMAGE_SIZE, CROP_SIZE)
    test_data = MVTecDataset(root=os.path.join(DATA_PATH, 'bottle'),
                             transform=data_transform, gt_transform=gt_transform, phase="test")
    test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

    # Build model
    print('\nBuilding model + loading trained weights...')
    model = build_model().to(DEVICE)
    model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=DEVICE))
    model.eval()

    # Pre-compute RoPE for fixed grid size (enables CUDAGraphs in reduce-overhead mode)
    patch_H = CROP_SIZE // 16  # 448 / 16 = 28
    patch_W = CROP_SIZE // 16
    model.encoder.rope_embed.precompute(patch_H, patch_W)
    print(f'  RoPE pre-computed for {patch_H}x{patch_W} grid')

    model = try_compile(model)

    dummy_input = torch.randn(BATCH_SIZE, 3, CROP_SIZE, CROP_SIZE, device=DEVICE)

    # Speed benchmark
    print('\n--- Speed ---')
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()

    speed = benchmark_speed(model, dummy_input)
    vram = torch.cuda.max_memory_allocated() / 1024**2
    print(f'  Peak VRAM: {vram:.0f} MB')

    # Accuracy benchmark
    print('\n--- Accuracy ---')
    wrapped = AutocastWrapper(model)
    results = evaluation_batch(wrapped, test_dataloader, DEVICE, max_ratio=0.01, resize_mask=256)
    auroc_sp, ap_sp, f1_sp, auroc_px, ap_px, f1_px, aupro_px = results

    # Summary
    print('\n' + '=' * 60)
    print('RESULTS')
    print('=' * 60)
    print(f'  Speed:    {speed:.2f} ms/image')
    print(f'  VRAM:     {vram:.0f} MB')
    print(f'  I-AUROC:  {auroc_sp:.4f}')
    print(f'  I-AP:     {ap_sp:.4f}')
    print(f'  I-F1:     {f1_sp:.4f}')
    print(f'  P-AUROC:  {auroc_px:.4f}')
    print(f'  P-AP:     {ap_px:.4f}')
    print(f'  P-F1:     {f1_px:.4f}')
    print(f'  P-AUPRO:  {aupro_px:.4f}')
    print('=' * 60)

    del model, wrapped
    torch.cuda.empty_cache()