Upload eval_zero.py with huggingface_hub

eval_zero.py

@@ -0,0 +1,220 @@
+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()
+
+