eval_zero.py

import os
import argparse
import random
import numpy as np
import torch
from torch.nn import functional as F
from tqdm import tqdm
from CLIP.clip import create_model
from CLIP.adapter import CLIPAD
from sklearn.metrics import roc_auc_score, average_precision_score
from dataset.continual import ImageDataset
import csv
import logging
from CoOp import PromptMaker
import json
from safetensors.torch import load_file

os.environ["TOKENIZERS_PARALLELISM"] = "false"

import warnings
warnings.filterwarnings("ignore")

def setup_seed(seed):
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

def get_logger(output_dir):
    # set log file
    log_file = f"{output_dir}/log.log"
    head = '%(asctime)-15s %(message)s'
    logging.basicConfig(filename=log_file,
                        format=head)
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    console = logging.StreamHandler()
    logging.getLogger('').addHandler(console)

    return logger

def main():
    parser = argparse.ArgumentParser(description='Evaluation')
    parser.add_argument('--model_name', type=str, default='ViT-L-14-336', help="ViT-B-16-plus-240, ViT-L-14-336")
    parser.add_argument('--pretrain', type=str, default='openai', help="laion400m, openai")
    parser.add_argument('--img_size', type=int, default=336)
    parser.add_argument("--features_list", type=int, nargs="+", default=[6, 12, 18, 24], help="features used")
    parser.add_argument('--seed', type=int, default=111)
    parser.add_argument('--gpu', type=str, default="0")
    parser.add_argument("--meta_file", type=str, default="meta_files/meta_mvtec.json")
    parser.add_argument("--n_learnable_token", type=int, default=8, help="number of learnable token")
    parser.add_argument("--adapter_ckpt", type=str, default="scenario2/30classes/adapters_sc2_task2.safetensors", help="adapter checkpoint path")
    parser.add_argument("--prompt_makder_ckpt", type=str, default="scenario2/30classes/prompt_maker_sc2.safetensors", help="prompt maker checkpoint path")
    parser.add_argument("--save_path", type=str, default="results_zero")
    parser.add_argument("--data_root", type=str, default="data/mvtec_anomaly_detection")

    args = parser.parse_args()

    setup_seed(args.seed)

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:{}".format(args.gpu) if use_cuda else "cpu")

    save_path = args.save_path
    if not os.path.isdir(save_path):
        os.makedirs(save_path)
    
    # for logging
    logger = get_logger(save_path)
    logger.info(args)
    
    # fixed feature extractor
    clip_model = create_model(model_name=args.model_name, img_size=args.img_size, device=device, pretrained=args.pretrain, require_pretrained=True)
    
    # prompt learner
    prompts = {
        "normal": [
        "This is an example of a normal object",
        "This is a typical appearance of the object",
        "This is what a normal object looks like",
        "A photo of a normal object",
        "This is not an anomaly",
        "This is an example of a standard object.",
        "This is the standard appearance of the object.",
        "This is what a standard object looks like.",
        "A photo of a standard object.",
        "This object meets standard characteristics."
    ],
        "abnormal": [
        "This is an example of an anomalous object",
        "This is not the typical appearance of the object",
        "This is what an anomaly looks like",
        "A photo of an anomalous object",
        "An anomaly detected in this object",
        "This is an example of an abnormal object.",
        "This is not the usual appearance of the object.",
        "This is what an abnormal object looks like.",
        "A photo of an abnormal object.",
        "An abnormality detected in this object."
    ]
    }

    clip_model.device = device
    clip_model.to(device)

    prompt_maker = PromptMaker(
        prompts=prompts,
        clip_model=clip_model,
        n_ctx= args.n_learnable_token,
        CSC = True,
        class_token_position=['end'],
    ).to(device)

    model = CLIPAD(clip_model=clip_model, features=args.features_list)
    model.to(device)
    model.eval()

    # load checkpoint
    adpater_state_dict = load_file(args.adapter_ckpt)
    model.adapters.load_state_dict(adpater_state_dict)
    logger.info(f"load adapter from {args.adapter_ckpt}")
    prompt_state_dict = load_file(args.prompt_makder_ckpt)
    prompt_maker.prompt_learner.load_state_dict(prompt_state_dict)
    logger.info(f"load prompt maker from {args.prompt_makder_ckpt}")

    kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {}

    prompt_maker.eval()
    model.eval()
    
    logging.info(f"start zero shot {args.meta_file} test")
    task_meta = json.load(open(args.meta_file, 'r'))

    class_name_list = list(task_meta["test"].keys())
    test_dataset_list = [ImageDataset(data_root=args.data_root, meta_file=task_meta, resize=args.img_size, mode="test", test_class=class_name) for class_name in class_name_list]
    test_loader_list = [torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, **kwargs) for test_dataset in test_dataset_list]

    with torch.cuda.amp.autocast(), torch.no_grad():
        # test all class
        seg_ap_list = []
        img_auc_list = []
        prompt_maker.eval()
        model.eval()
        text_features = prompt_maker()

        for test_loader, class_name in zip(test_loader_list, class_name_list):
            logger.info(f"start test {class_name}")
            roc_auc_im, seg_ap = test(args, model, test_loader, text_features, device)
            logger.info(f'{class_name} P-AP : {round(seg_ap,4)}')
            logger.info(f'{class_name} I-AUC : {round(roc_auc_im, 4)}')
            seg_ap_list.append(seg_ap)
            img_auc_list.append(roc_auc_im)

        seg_ap_mean = np.mean(seg_ap_list)
        img_auc_mean = np.mean(img_auc_list)

        logger.info(f'Average P-AP : {round(seg_ap_mean,4)}')
        logger.info(f'Average I-AUC : {round(img_auc_mean, 4)}')


def test(args, model, test_loader, text_features, device):
    gt_list = []
    gt_mask_list = []

    seg_score_map_zero = []
    image_scores = []
    for data in tqdm(test_loader):
        image, mask, cls_name, label = data['image'], data['mask'], data['cls_name'], data['anomaly']
        image = image.to(device)
        mask[mask > 0.5], mask[mask <= 0.5] = 1, 0

        with torch.no_grad(), torch.cuda.amp.autocast():
            _, ada_patch_tokens = model(image)
            ada_patch_tokens = [p[0, 1:, :] for p in ada_patch_tokens]

            anomaly_maps = []
            image_score = 0
            for layer in range(len(ada_patch_tokens)):
                ada_patch_tokens[layer] /= ada_patch_tokens[layer].norm(dim=-1, keepdim=True)
                anomaly_map = (100.0 * ada_patch_tokens[layer] @ text_features).unsqueeze(0)
                B, L, C = anomaly_map.shape
                H = int(np.sqrt(L))

                # image
                anomaly_score = torch.softmax(anomaly_map, dim=-1)[:, :, 1]
                image_score += anomaly_score.max()
        
                anomaly_maps.append(anomaly_map)

            score_map = torch.mean(torch.stack(anomaly_maps, dim=1), dim=1)
            score_map = F.interpolate(score_map.permute(0, 2, 1).view(B, 2, H, H),
                                        size=args.img_size, mode='bilinear', align_corners=True)
            score_map = torch.softmax(score_map, dim=1)[:, 1, :, :]
            score_map = score_map.squeeze(0).cpu().numpy()
            seg_score_map_zero.append(score_map)
            image_scores.append(image_score.cpu() / len(ada_patch_tokens))
                        
            gt_mask_list.append(mask.squeeze().cpu().detach().numpy())
            gt_list.extend(label.cpu().detach().numpy())

            
    gt_list = np.array(gt_list)
    gt_mask_list = np.asarray(gt_mask_list)
    gt_mask_list = (gt_mask_list>0).astype(np.int_)

    segment_scores = np.array(seg_score_map_zero)
    image_scores = np.array(image_scores)

    roc_auc_im = roc_auc_score(gt_list, image_scores)

    seg_pr = average_precision_score(gt_mask_list.flatten(), segment_scores.flatten())

    return roc_auc_im, seg_pr


if __name__ == '__main__':
    main()