DINOv3+MobileCLIP2高效协同，LiZAD轻量级零样本异常检测赋能工业智能制造，推理实测30ms以内，内存仅1.01 GB

conda create -n LiZAD python=3.12

conda activate LiZAD

pip install -r requirements.txt

git clone https://github.com/apple/ml-mobileclip.git
cd ml-mobileclip
pip install .

import os
import csv
import time
from argparse import ArgumentParser

import cv2
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm

from utils.args import Args
from datasets.constants import DatasetConstants
from utils.utils import generate_text_embeddings
from utils.transformations import get_transforms
from backbones.encoders import ImageEncoder
from model.model import ZSADModel

IMAGE_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.bmp', '.tif')


def collect_images(input_path):
    if os.path.isfile(input_path):
        return [input_path]
    if os.path.isdir(input_path):
        img_paths = []
        for root, dirs, files in os.walk(input_path):
            dirs[:] = [d for d in dirs if d.lower() != 'ground_truth']  # skip gt masks in dataset layouts
            for f in files:
                if f.lower().endswith(IMAGE_EXTENSIONS):
                    img_paths.append(os.path.join(root, f))
        return sorted(img_paths)
    raise FileNotFoundError(f"Input path does not exist: {input_path}")


def infer_cls_name(img_path, class_names):
    parts = [p.lower() for p in os.path.normpath(img_path).split(os.sep)]
    for cls in class_names:
        if cls.lower() in parts:
            return cls
    return None


def load_image(img_path):
    if img_path.lower().endswith('.tif'):
        img = cv2.imread(img_path)
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    else:
        img = Image.open(img_path).convert('RGB')
    return img


def save_heatmap_overlay(img_path, anomaly_map, save_path):
    img = cv2.imread(img_path)
    h, w = img.shape[:2]

    heatmap = cv2.resize(anomaly_map, (w, h), interpolation=cv2.INTER_LINEAR)
    heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)

    overlay = cv2.addWeighted(img, 0.5, heatmap, 0.5, 0)

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    cv2.imwrite(save_path, overlay)


def predict_batch(batch, model, image_encoder, text_embeddings, device):
    imgs = torch.stack([b['img'] for b in batch], dim=0)

    normal_batch = torch.stack([text_embeddings['normal'][b['cls_name']] for b in batch], dim=0)
    abnormal_batch = torch.stack([text_embeddings['abnormal'][b['cls_name']] for b in batch], dim=0)

    text_embeddings_dict = {
        'normal': normal_batch.to(device),
        'abnormal': abnormal_batch.to(device)
    }

    cls, patches = image_encoder(imgs)

    text_guided_anomaly_scores, global_context_regularizer_score = model(text_embeddings_dict, [cls, patches])

    pixel_anomaly_maps = text_guided_anomaly_scores[:, 1, :, :].cpu().detach().numpy().astype(np.float32)
    image_scores = torch.softmax(global_context_regularizer_score, dim=1)[:, 1].cpu().detach().numpy()

    return pixel_anomaly_maps, image_scores


def predict(cli_args):
    dataset_constants = DatasetConstants('./', cli_args.dataset_name)
    class_names = dataset_constants.get_class_names()

    img_paths = collect_images(cli_args.input)
    if len(img_paths) == 0:
        raise FileNotFoundError(f"No images found in: {cli_args.input}")

    # Resolve class name for every image before loading any model
    cls_names = []
    for img_path in img_paths:
        cls_name = cli_args.cls_name or infer_cls_name(img_path, class_names)
        if cls_name is None:
            raise ValueError(
                f"Could not infer class name from path: {img_path}\n"
                f"Please specify it explicitly with --cls_name (one of: {class_names})"
            )
        cls_names.append(cls_name)

    args = Args(
        vision_model_id="facebook/dinov3-vits16-pretrain-lvd1689m",
        model_id="MobileCLIP2-S0",
        vision_layers=[3, 5, 7, 11],
        text_dim=512,
        vision_dim=384,
        base_dir="./",
        batch_size=cli_args.batch_size,
        img_size=518,
        lr=1e-4,
        out_dim=256,
        device=cli_args.device,
    )

    print(f"Loading text embeddings for dataset: {cli_args.dataset_name}")
    text_embeddings = generate_text_embeddings(args, dataset_constants)

    transform_img, _ = get_transforms(args.img_size)

    model = ZSADModel(args).to(args.device)
    print(f"Loading model: {cli_args.model_path}")
    checkpoint = torch.load(cli_args.model_path, map_location=args.device)
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)
    model.eval()

    image_encoder = ImageEncoder(args.vision_model_id, args.vision_layers, device=args.device)

    os.makedirs(cli_args.output_dir, exist_ok=True)
    input_is_dir = os.path.isdir(cli_args.input)

    results = []
    with torch.no_grad():
        for start in tqdm(range(0, len(img_paths), args.batch_size), desc="Predicting"):
            batch_paths = img_paths[start:start + args.batch_size]
            batch = []
            for img_path, cls_name in zip(batch_paths, cls_names[start:start + args.batch_size]):
                batch.append({
                    'img': transform_img(load_image(img_path)),
                    'cls_name': cls_name,
                    'img_path': img_path,
                })

            if args.device.startswith('cuda'):
                torch.cuda.synchronize()
            t0 = time.perf_counter()
            pixel_anomaly_maps, image_scores = predict_batch(batch, model, image_encoder, text_embeddings, args.device)
            if args.device.startswith('cuda'):
                torch.cuda.synchronize()
            per_image_time_ms = (time.perf_counter() - t0) / len(batch) * 1000

            for item, anomaly_map, image_score in zip(batch, pixel_anomaly_maps, image_scores):
                if input_is_dir:
                    rel = os.path.relpath(item['img_path'], cli_args.input)
                else:
                    rel = os.path.basename(item['img_path'])
                save_path = os.path.join(cli_args.output_dir, os.path.splitext(rel)[0] + '_heatmap.png')

                save_heatmap_overlay(item['img_path'], anomaly_map, save_path)

                results.append({
                    'img_path': item['img_path'],
                    'cls_name': item['cls_name'],
                    'anomaly_score': float(image_score),
                    'max_pixel_score': float(anomaly_map.max()),
                    'inference_time_ms': per_image_time_ms,
                    'heatmap_path': save_path,
                })

    csv_path = os.path.join(cli_args.output_dir, 'results.csv')
    with open(csv_path, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['img_path', 'cls_name', 'anomaly_score', 'max_pixel_score', 'inference_time_ms', 'heatmap_path'])
        writer.writeheader()
        writer.writerows(results)

    for r in results:
        print(f"{r['img_path']}  cls={r['cls_name']}  anomaly_score={r['anomaly_score']:.4f}  max_pixel_score={r['max_pixel_score']:.4f}  inference_time={r['inference_time_ms']:.1f}ms")

    avg_time_ms = sum(r['inference_time_ms'] for r in results) / len(results)
    print(f"\nDone. {len(results)} image(s) predicted, avg inference time: {avg_time_ms:.1f} ms/image. Heatmaps and results.csv saved to: {cli_args.output_dir}")