eval_continual.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/scenario1_5classes_tasks.json")
    parser.add_argument("--base_meta_file", type=str, default="meta_files/scenario1_base.json")    
    parser.add_argument("--num_tasks", type=int, default=12, help="number of tasks")
    parser.add_argument("--n_learnable_token", type=int, default=8, help="number of learnable token")
    parser.add_argument("--checkpoints", type=str, default="scenario2/30classes", help="folder path to checkpoints")
    parser.add_argument("--checkpoint_base", type=str, default="scenario2/adapters_base.safetensors", help="checkpoint base path")
    parser.add_argument("--prompt_makder_ckpt", type=str, default="scenario2/prompt_maker.safetensors", help="prompt maker checkpoint path")
    parser.add_argument("--task_id", type=int, default=1, help="test task id")  # 0 - base classes
    parser.add_argument("--save_path", type=str, default="results")
    parser.add_argument("--data_root", type=str, default="data")

    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
    if args.task_id == 0:
        checkpoint_path = args.checkpoint_base
    else:
        checkpoint_path = f"{args.checkpoints}/adapters_task{args.task_id}.safetensors"

    checkpoint = load_file(checkpoint_path)
    model.adapters.load_state_dict(checkpoint)
    logger.info(f"load adapter from {checkpoint_path}")
    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 {}

    # save results
    num_tasks = args.num_tasks + 1
    results_image = np.full((num_tasks, num_tasks), 0)  # save for csv image-level
    results_pixel = np.full((num_tasks, num_tasks), 0)  # save for csv pixel-level

    # load saved_results
    csv_image = f"{save_path}/results_image.csv"
    csv_pixel = f"{save_path}/results_pixel.csv"
    
    if os.path.exists(csv_image):
        with open(csv_image, mode="r") as file:
            reader = csv.reader(file)
            for i, row in enumerate(reader):
                if not i == 0:
                    results_image[i-1] = row
        logger.info(f"load previous results from {csv_image}")

    if os.path.exists(csv_pixel):
        with open(csv_pixel, mode="r") as file:
            reader = csv.reader(file)
            for i, row in enumerate(reader):
                if not i == 0:
                    results_pixel[i-1] = row
        logger.info(f"load previous results from {csv_pixel}")

    prompt_maker.eval()
    model.eval()
    
    task_all_meta_info = json.load(open(args.meta_file, 'r'))

    # test all previous tasks
    for i in range(args.task_id + 1):
        if i == 0: # base classes
            task_meta = json.load(open(args.base_meta_file, 'r'))
            logging.info(f"start base task test")
        else:
            task_meta = task_all_meta_info[f"task_{i}"]
            logging.info(f"start task_{i} test")

        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)}')

        # save results csv (i task)
        seg_ap_mean = round(seg_ap_mean,4)
        img_auc_mean = round(img_auc_mean, 4)
        results_image[args.task_id, i] = img_auc_mean
        results_pixel[args.task_id, i] = seg_ap_mean

        logger.info(f"save results csv task {i}")

        # save results csv
        with open(csv_image, mode="w", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(["Base"] + ["Task " + str(i + 1) for i in range(num_tasks-1)])
            for row in results_image:
                writer.writerow(row)
        with open(csv_pixel, mode="w", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(["Base"] + ["Task " + str(i + 1) for i in range(num_tasks-1)])
            for row in results_pixel:
                writer.writerow(row)

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()