test_mvlogo.py

# from models.RB_VIT import RB_VIT
import torch
from config import DefaultConfig
import os
from torch import optim

import torch.nn.functional as F
from tqdm import tqdm

from matplotlib import pyplot as plt
import numpy as np
import argparse
from sklearn.metrics import roc_auc_score, f1_score, average_precision_score
from utils.metric import cal_pro_metric_new
import cv2
from models.misc import NativeScalerWithGradNormCount as NativeScaler
from losses.vgg_loss import StyleLoss, PerceptualLoss
from scipy.ndimage import gaussian_filter
from datasets.dataset import denormalize
from ref_find import get_pos_sample
from models.TFA_Net_model import *
class Model(object):
    def __init__(self, opt, test_no_mask, test_add):
        super(Model, self).__init__()
        # if test_no_mask == True:
        #     model_name = 'RB_VIT_dir_res_ref'
        # else:
        #     model_name = opt.model_name
        self.opt = opt
        self.model = eval(opt.model_name)(opt)
        self.device = opt.device
        self.test_add = test_add
        self.class_name = opt.class_name
        self.trainloader = opt.trainloader
        self.testloader = opt.testloader
        self.loss_scaler = NativeScaler()

        if self.opt.resume != "":
            print('\nload pre-trained networks')
            self.opt.iter = torch.load(os.path.join(self.opt.resume, f'{opt.model_name}_{opt.backbone_name}_k={opt.k}.pth'))['epoch']
            print(self.opt.iter)
            self.model.load_state_dict(torch.load(os.path.join(self.opt.resume, f'{opt.model_name}_{opt.backbone_name}_k={opt.k}.pth'))['state_dict'], strict=False)
            print('\ndone.\n')

        if self.opt.isTrain:
            self.model.Roncon_model.train()
            self.optimizer_g = optim.AdamW(self.model.Roncon_model.parameters(), lr=opt.lr, betas=(0.9, 0.95))
        if test_no_mask == True:
            self.save_root = f"./result_logo_test/{opt.model_name}_{opt.backbone_name}_test_no_mask"
        else:
            self.save_root = f"./result_test/{opt.model_name}_{opt.backbone_name}"

        if test_add == True:
            self.save_root += 'add'
        # os.makedirs(os.path.join(self.save_root, "weight"), exist_ok=True)
        self.ckpt_root = os.path.join(self.save_root, "weight/{}".format(self.class_name))
        self.vis_root = os.path.join(self.save_root, "img/{}".format(self.class_name))



    def get_max(self, tensor):
        a_1, _ = torch.max(tensor, dim=1, keepdim=True)
        a_2, _ = torch.max(a_1, dim=2, keepdim=True)
        a_3, _ = torch.max(a_2, dim=3, keepdim=True)
        return a_3
    def train(self):

        loss_now = 100000
        auc_now = 0.
        for epoch in range(self.opt.iter, self.opt.niter):
            self.model.Feature_extractor.eval()
            self.model.Roncon_model.train(True)
            self.model.to(self.device)
            loss_total = 0.
            count = 0
            for index, (x, _, _, _) in enumerate(tqdm(self.trainloader, ncols=80)):
                bs = x.shape[0]
                x = x.to(self.device)
                ref_x = get_pos_sample(self.opt.referenc_img_file, self.device, bs)

                deep_feature,_, recon_feature, loss = self.model(x, ref_x)
                self.loss_scaler(loss, self.optimizer_g, parameters=self.model.Roncon_model.parameters(), update_grad=(index + 1) % 1 == 0)
                loss_total += loss.item()
                count += 1

            loss_total = loss_total / count
            print('the {} epoch is done   loss:{}'.format(epoch + 1, loss_total))
            if (epoch + 1) % 50 == 0:
                 # self.test_2()
                x1, x2, x3, x4 = self.test()
                auc_roc = x1+x2
                if auc_roc > auc_now:
                    auc_now = auc_roc
                    class_rocauc[self.opt.class_name] = (x1, x2, x3, x4)
                    print('save model')
                    weight_dir = self.ckpt_root
                    os.makedirs(weight_dir, exist_ok=True)
                    torch.save({'epoch': epoch + 1, 'state_dict': self.model.state_dict()},
                               f'%s/{self.opt.model_name}_{self.opt.backbone_name}.pth' % (weight_dir))



    def cal_auc(self, score_list, logo_score_list, struct_score_list, score_map_list, test_y_list, logo_y_list, struct_y_list, test_mask_list):
        flatten_y_list = np.array(test_y_list).ravel()
        flatten_score_list = np.array(score_list).ravel()
        image_level_ROCAUC = roc_auc_score(flatten_y_list, flatten_score_list)
        flatten_logo_y_list = np.array(logo_y_list).ravel()
        flatten_logo_score_list = np.array(logo_score_list).ravel()
        logo_img_auroc = roc_auc_score(flatten_logo_y_list, flatten_logo_score_list)
        flatten_stru_y_list = np.array(struct_y_list).ravel()
        flatten_stru_score_list = np.array(struct_score_list).ravel()
        stru_img_auroc = roc_auc_score(flatten_stru_y_list, flatten_stru_score_list)
        # image_level_AP = average_precision_score(flatten_y_list, flatten_score_list)

        flatten_mask_list = np.concatenate(test_mask_list).ravel()
        flatten_score_map_list = np.concatenate(score_map_list).ravel()
        pixel_level_ROCAUC = roc_auc_score(flatten_mask_list, flatten_score_map_list)
        # pixel_level_AP = average_precision_score(flatten_mask_list, flatten_score_map_list)
        # pro_auc_score = 0
        # pro_auc_score = cal_pro_metric_new(test_mask_list, score_map_list, fpr_thresh=0.3)
        return round(image_level_ROCAUC,3), round(pixel_level_ROCAUC,3), round(logo_img_auroc,3), round(stru_img_auroc,3)
        # return  image_level_ROCAUC, pixel_level_ROCAUC

    def F1_score(self, score_map_list, test_mask_list):
        flatten_mask_list = np.concatenate(test_mask_list).ravel()
        flatten_score_map_list = np.concatenate(score_map_list).ravel()
        F1_score = f1_score(flatten_mask_list, flatten_score_map_list)
        return F1_score

    def filter(self, pred_mask):
        pred_mask_my = np.squeeze(np.squeeze(pred_mask, 0), 0)
        pred_mask_my = cv2.medianBlur(np.uint8(pred_mask_my * 255), 7)
        mean = np.mean(pred_mask_my)
        std = np.std(pred_mask_my)
        _ , binary_pred_mask = cv2.threshold(pred_mask_my, mean+2.75*std, 255, type=cv2.THRESH_BINARY)
        binary_pred_mask = np.uint8(binary_pred_mask/255)
        pred_mask_my = np.expand_dims(np.expand_dims(pred_mask_my, 0), 0)
        binary_pred_mask = np.expand_dims(np.expand_dims(binary_pred_mask, 0), 0)
        return pred_mask_my, binary_pred_mask


    # def thresholding(self, pred_mask_my):
    #     np_img

        # return
    # def feature_map_vis(self, feature_map_list):
    #     # feature_map_list = [torch.mean(i.clone(), dim=1).squeeze(0).cpu().detach().numpy() for i in feature_map_list]
    #     feature_map_list = [(i.squeeze(0))[25, :, :].cpu().detach().numpy() for i in feature_map_list]
    #     return feature_map_list
    def feature_map_vis(self, feature_map_list):
        feature_map_list = [torch.mean(i.clone(), dim=1).squeeze(0).cpu().detach().numpy() for i in feature_map_list]
        return feature_map_list

    def test(self):
        test_y_list = []
        logo_y_list=[]
        struct_y_list = []
        test_mask_list = []
        score_list = []
        score_map_list = []
        logo_score_list = []
        struct_score_list = []

        for idx, (x, y, mask, name) in enumerate(tqdm(self.testloader, ncols=80)):
            test_y_list.extend(y.detach().cpu().numpy())
            test_mask_list.extend(mask.detach().cpu().numpy())

            self.model.eval()
            self.model.to(self.device)
            x = x.to(self.device)
            mask = mask.to(self.device)
            mask_cpu = mask.cpu().detach().numpy()[0, :, :, :].transpose((1, 2, 0))
            ref_x = get_pos_sample(self.opt.referenc_img_file, self.device, 1)
            deep_feature, ref_feature, recon_feature, _ = self.model(x, ref_x, None)
            feature_map_vis_list = self.feature_map_vis([deep_feature, ref_feature, recon_feature])
            dis_amap, dir_amap = self.model.a_map(deep_feature, recon_feature)
            dis_amap = gaussian_filter(dis_amap, sigma=4)
            dir_amap = gaussian_filter(dir_amap, sigma=4)
            # print(type(name0]))
            name_list= name[0].split(r'!')
            # print(name_list)
            category, img_name = name_list[-2], name_list[-1]
            if category == 'logical_anomalies':
                logo_y_list.extend(y.detach().cpu().numpy())
            elif category == 'structural_anomalies':
                struct_y_list.extend(y.detach().cpu().numpy())
            if self.test_add == False:
                amap = dir_amap*dis_amap
                # amap = dir_amap*5+dis_amap
            else:
                # print('ok')
                # print(np.max(dir_amap))
                # print(np.max(dis_amap))
                amap = 0.5*(dir_amap/np.max(dir_amap)) + 0.5*(dis_amap/np.max(dis_amap))
            self.save_img([x,*feature_map_vis_list, dis_amap, dir_amap, amap, mask_cpu], os.path.join(self.vis_root, category), img_name)


            score_list.extend(np.array(np.std(amap)).reshape(1))
            score_map_list.extend(amap.reshape((1, 1, 256, 256)))
            if category == 'logical_anomalies':
                logo_score_list.extend(np.array(np.std(amap)).reshape(1))
            elif category == 'structural_anomalies':
                struct_score_list.extend(np.array(np.std(amap)).reshape(1))

        image_level_ROCAUC, pixel_level_ROCAUC, image_level_AP, pixel_level_AP= self.cal_auc(score_list, logo_score_list, struct_score_list, score_map_list, test_y_list, logo_y_list, struct_y_list, test_mask_list
                                                                                             )
        # F1_score = self.F1_score(F1_score_map_list, test_mask_list)
        print('image_auc_roc: {} '.format(image_level_ROCAUC),
              'pixel_auc_roc: {} '.format(pixel_level_ROCAUC),
              'logi_auroc: {}'.format(image_level_AP),
              'stru_auroc: {}'.format(pixel_level_AP)
             )
        class_rocauc[self.opt.class_name] = (image_level_ROCAUC, pixel_level_ROCAUC, image_level_AP, pixel_level_AP)
        return image_level_ROCAUC, pixel_level_ROCAUC, image_level_AP, pixel_level_AP

    def vis_img(self, img_list, save_root, idx_name):
        os.makedirs(save_root, exist_ok=True)
        input_frame = denormalize(img_list[0].clone().squeeze(0).cpu().detach().numpy())
        cv2_input = np.array(input_frame, dtype=np.uint8)
        plt.figure()
        plt.subplot(241)
        plt.imshow(cv2_input)
        plt.axis('off')
        plt.subplot(242)
        plt.imshow(img_list[1])
        plt.axis('off')
        plt.subplot(243)
        plt.imshow(img_list[2])
        plt.axis('off')
        plt.subplot(244)
        plt.imshow(img_list[3])
        plt.axis('off')
        plt.subplot(245)
        plt.imshow(img_list[4], cmap='jet')
        plt.axis('off')
        plt.subplot(246)
        plt.imshow(img_list[5],cmap='jet')
        plt.axis('off')
        plt.subplot(247)
        plt.imshow(img_list[6], cmap='jet')
        plt.axis('off')
        plt.subplot(248)
        plt.imshow(img_list[7])
        plt.axis('off')
        plt.savefig(os.path.join(save_root, idx_name))
        plt.close()

    def save_img(self, img_list, save_root, idx_name):
        os.makedirs(save_root, exist_ok=True)
        input_frame = denormalize(img_list[0].clone().squeeze(0).cpu().detach().numpy())
        cv2_input = np.array(input_frame, dtype=np.uint8)
        # plt.figure()
        # plt.subplot(241)
        plt.imsave(os.path.join(save_root, f'{idx_name}_{0}.png'), cv2_input)
        plt.imsave(os.path.join(save_root, f'{idx_name}_{1}.png'), img_list[1])
        plt.imsave(os.path.join(save_root, f'{idx_name}_{2}.png'), img_list[2])
        plt.imsave(os.path.join(save_root, f'{idx_name}_{3}.png'), img_list[3])
        plt.imsave(os.path.join(save_root, f'{idx_name}_{4}.png'), img_list[4], cmap='jet')
        plt.imsave(os.path.join(save_root, f'{idx_name}_{5}.png'), img_list[5], cmap='jet')
        plt.imsave(os.path.join(save_root, f'{idx_name}_{6}.png'), img_list[6], cmap='jet')

        plt.imsave(os.path.join(save_root, f'{idx_name}_{7}.png'),cv2.cvtColor(img_list[7], cv2.COLOR_GRAY2RGB), cmap='gray')
        # plt.axis('off')
        # plt.subplot(242)
        # plt.imwrite()
        # # plt.axis('off')
        # # plt.subplot(243)
        # plt.imshow(img_list[2])
        # # plt.axis('off')
        # # plt.subplot(244)
        # plt.imshow(img_list[3])
        # # plt.axis('off')
        # # plt.subplot(245)
        # plt.imshow(img_list[4], cmap='jet')
        #
        # # plt.axis('off')
        # # plt.subplot(246)
        # plt.imshow(img_list[5],cmap='jet')
        # # plt.axis('off')
        # # plt.subplot(247)
        # plt.imshow(img_list[6], cmap='jet')
        # # plt.axis('off')
        # # plt.subplot(248)
        # plt.imshow(img_list[7])
        # plt.axis('off')
        # plt.savefig(os.path.join(save_root, idx_name))
        # plt.close()

    def tensor_to_np_cpu(self, tensor):
        x_cpu = tensor.squeeze(0).data.cpu().numpy()
        x_cpu = np.transpose(x_cpu, (1, 2, 0))
        return x_cpu

    def check(self, img):
        if len(img.shape) == 2:
            return img
        if img.shape[2] == 3:
            return img
        elif img.shape[2] == 1:
            return img.reshape(img.shape[0], img.shape[1])

MVTec_CLASS_NAMES = [ 'breakfast_box', 'juice_bottle', 'pushpins', 'screw_bag', 'splicing_connectors']

class_rocauc = {
                'breakfast_box':(0, 0, 0, 0),
                'juice_bottle':(0, 0, 0, 0),
                'pushpins':(0, 0, 0, 0),
                'screw_bag':(0, 0, 0, 0),
                'splicing_connectors':(0, 0, 0, 0)
            }

model_name_list = ['RB_VIT', 'RB_VIT_dir', 'VIT_dir', 'RB_VIT_average', 'RB_VIT_dir_mask25', 'RB_VIT_dir_mask50', 'RB_VIT_dir_mask75', 'ST_VIT', 'RB_VIT_Res', 'RB_VIT_dir_no_ref',
                   'RB_VIT_dir_res_ref', 'RB_VIT_dir_res_ref_mask25', 'RB_VIT_dir_res_ref_mask50', 'RB_VIT_dir_res_ref_mask75',
                   'RB_VIT_dir_res_ref_mask_random']

if __name__ == '__main__':
    opt = DefaultConfig()
    test_no_mask = True
    test_add = False
    from datasets.logodataset import MVTecDataset
    from torch.utils.data import DataLoader
    opt.model_name = model_name_list[1]
    for classname in MVTec_CLASS_NAMES:
        opt.class_name = classname
        # opt.class_name = 'capsule'
        opt.referenc_img_file =  f'data/mvtec_anomaly_detection/{opt.class_name}/train/good/000.png'
        save_name = opt.model_name+'_'+opt.backbone_name
        opt.resume = fr'F:\LW\RB-VIT\result/{save_name}/weight/{opt.class_name}'
        # if test_no_mask == True:
        #     opt.model_name = 'RB_VIT_dir_res_ref'
        print(opt.class_name, opt.model_name)
        # print(opt.referenc_img_file)
        # opt.resume = r'result/RB_VIT_dir_res_ref_VGG/weight/capsule'
        opt.train_dataset = MVTecDataset(dataset_path=opt.data_root, class_name=opt.class_name, is_train=True)
        opt.test_dataset = MVTecDataset(dataset_path=opt.data_root, class_name=opt.class_name, is_train=False)
        opt.trainloader = DataLoader(opt.train_dataset, batch_size=opt.batch_size, shuffle=True)
        opt.testloader = DataLoader(opt.test_dataset, batch_size=1, shuffle=False)
        model = Model(opt, test_no_mask, test_add)
        model.test()
    print(class_rocauc)
    value = list(class_rocauc.values())
    img_roc = [i[0] for i in value]
    pixel_roc = [i[1] for i in value]
    img_ap = [i[2] for i in value]
    pixel_ap = [i[3] for i in value]
    mean_img_roc = np.mean(np.array(img_roc))
    mean_pixel_roc = np.mean(np.array(pixel_roc))
    mean_img_ap = np.mean(np.array(img_ap))
    mean_pixel_ap = np.mean(np.array(pixel_ap))

    print(round(mean_img_roc,3), round(mean_pixel_roc,3), round(mean_img_ap, 3), round(mean_pixel_ap,3))