This is an unofficial implementation of the paper "Towards Total Recall in Industrial Anomaly Detection".
import os
import numpy as np
from utils.config import ConfigData
from datasets.shared_memory import SharedMemory


class MVTecDataset:
    @classmethod
    def __init__(cls, type_data):
        cls.type_data = type_data
        cls.SHAPE_INPUT = ConfigData.SHAPE_INPUT

        cls.imgs_train = None
        cls.imgs_test = {}
        cls.gts_test = {}

        # read train data
        desc = 'read images for train (case:good)'
        path = os.path.join(ConfigData.path_parent, type_data, 'train/good')
        files = [os.path.join(path, f) for f in os.listdir(path)
                 if (os.path.isfile(os.path.join(path, f)) & ('.png' in f))]
        cls.files_train = np.sort(np.array(files))
        cls.imgs_train = SharedMemory.read_img_parallel(files=cls.files_train,
                                                        imgs=cls.imgs_train,
                                                        desc=desc)
        if ConfigData.shuffle:
            np.random.seed(ConfigData.seed)
            cls.imgs_train = np.random.permutation(cls.imgs_train)
        if ConfigData.flip_horz:
            cls.imgs_train = np.concatenate([cls.imgs_train,
                                             cls.imgs_train[:, :, ::-1]], axis=0)
        if ConfigData.flip_vert:
            cls.imgs_train = np.concatenate([cls.imgs_train,
                                             cls.imgs_train[:, ::-1]], axis=0)

        # read test data
        cls.files_test = {}
        cls.types_test = os.listdir(os.path.join(ConfigData.path_parent, type_data, 'test'))
        cls.types_test = np.array(sorted(cls.types_test))
        for type_test in cls.types_test:
            desc = 'read images for test (case:%s)' % type_test
            path = os.path.join(ConfigData.path_parent, type_data, 'test', type_test)
            files = [os.path.join(path, f) for f in os.listdir(path)
                     if (os.path.isfile(os.path.join(path, f)) & ('.png' in f))]
            cls.files_test[type_test] = np.sort(np.array(files))
            cls.imgs_test[type_test] = None
            cls.imgs_test[type_test] = SharedMemory.read_img_parallel(files=cls.files_test[type_test],
                                                                      imgs=cls.imgs_test[type_test],
                                                                      desc=desc)

        # read ground truth of test data
        for type_test in cls.types_test:
            # create memory shared variable
            if (type_test == 'good'):
                cls.gts_test[type_test] = np.zeros([len(cls.files_test[type_test]),
                                                    ConfigData.SHAPE_INPUT[0],
                                                    ConfigData.SHAPE_INPUT[1]], dtype=np.uint8)
            else:
                desc = 'read ground-truths for test (case:%s)' % type_test
                cls.gts_test[type_test] = None
                cls.gts_test[type_test] = SharedMemory.read_img_parallel(files=cls.files_test[type_test],
                                                                         imgs=cls.gts_test[type_test],
                                                                         is_gt=True,
                                                                         desc=desc)



datasets/mvtec_dataset.py
any-tech/PatchCore-ex
import numpy as np
import cv2
from tqdm import tqdm
import multiprocessing as mp
from multiprocessing.sharedctypes import RawArray
from utils.config import ConfigData


def read_and_resize(file):
    img = cv2.imread(file)[..., ::-1]  # BGR2RGB
    img = cv2.resize(img, (ConfigData.SHAPE_MIDDLE[1], ConfigData.SHAPE_MIDDLE[0]),
                     interpolation=cv2.INTER_AREA)
    img = img[ConfigData.pixel_cut[0]:(ConfigData.SHAPE_INPUT[0] + ConfigData.pixel_cut[0]),
              ConfigData.pixel_cut[1]:(ConfigData.SHAPE_INPUT[1] + ConfigData.pixel_cut[1])]
    SharedMemory.shared_array[np.where(SharedMemory.files == file)[0]] = img


def read_and_resize_ground_truth(file):
    file_gt = file.replace('/test/', '/ground_truth/')
    file_gt = file_gt.replace('.png', '_mask.png')

    gt = cv2.imread(file_gt, cv2.IMREAD_GRAYSCALE)
    gt = cv2.resize(gt, (ConfigData.SHAPE_MIDDLE[1], ConfigData.SHAPE_MIDDLE[0]),
                    interpolation=cv2.INTER_NEAREST)
    gt = gt[ConfigData.pixel_cut[0]:(ConfigData.SHAPE_INPUT[0] + ConfigData.pixel_cut[0]),
            ConfigData.pixel_cut[1]:(ConfigData.SHAPE_INPUT[1] + ConfigData.pixel_cut[1])]

    if np.max(gt) != 0:
        gt = (gt / np.max(gt)).astype(np.uint8)

    SharedMemory.shared_array[np.where(SharedMemory.files == file)[0]] = gt


class SharedMemory:
    @classmethod
    def read_img_parallel(cls, files, imgs, is_gt=False, desc='read images'):
        cls.files = files
        cls.shared_array = imgs

        if not is_gt:
            shape = (len(cls.files), ConfigData.SHAPE_INPUT[0], ConfigData.SHAPE_INPUT[1], 3)
            num_elm = shape[0] * shape[1] * shape[2] * shape[3]
        else:
            shape = (len(cls.files), ConfigData.SHAPE_INPUT[0], ConfigData.SHAPE_INPUT[1])
            num_elm = shape[0] * shape[1] * shape[2]

        ctype = np.ctypeslib.as_ctypes_type(np.dtype(np.uint8))
        data = np.ctypeslib.as_array(RawArray(ctype, num_elm))
        data.shape = shape
        cls.shared_array = data.view(np.uint8)

        # exec imread and imresize on multiprocess
        mp.set_start_method('fork', force=True)
        p = mp.Pool(min(mp.cpu_count(), ConfigData.num_cpu_max))

        func = read_and_resize
        if is_gt:
            func = read_and_resize_ground_truth

        for _ in tqdm(p.imap_unordered(func, cls.files), total=len(cls.files), desc=desc):
            pass

        p.close()

        return cls.shared_array

datasets/shared_memory.py
any-tech/PatchCore-ex
import os
import numpy as np
from argparse import ArgumentParser

from utils.config import ConfigData, ConfigFeat, ConfigPatchCore, ConfigDraw
from utils.tictoc import tic, toc
from utils.metrics import calc_imagewise_metrics, calc_pixelwise_metrics
from utils.visualize import draw_roc_curve, draw_distance_graph, draw_heatmap
from datasets.mvtec_dataset import MVTecDataset
from models.feat_extract import FeatExtract
from models.patchcore import PatchCore


def arg_parser():
    parser = ArgumentParser()
    # environment related
    parser.add_argument('-n', '--num_cpu_max', default=4, type=int,
                        help='number of CPUs for parallel reading input images')
    parser.add_argument('-c', '--cpu', action='store_true', help='use cpu')
    # I/O and visualization related
    parser.add_argument('-pp', '--path_parent', type=str, default='./mvtec_anomaly_detection',
                        help='parent path of data input path')
    parser.add_argument('-pr', '--path_result', type=str, default='./result',
                        help='output path of figure image as the evaluation result')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='save visualization of localization')
    parser.add_argument('-srf', '--size_receptive_field', type=int, default=15,
                        help='estimate and specify receptive field size (odd number)')
    parser.add_argument('-mv', '--mode_visualize', type=str, default='eval',
                        choices=['eval', 'infer'], help='set mode, [eval] or [infer]')
    # data loader related
    parser.add_argument('-bs', '--batch_size', type=int, default=16,
                        help='batch-size for feature extraction by ImageNet model')
    parser.add_argument('-sr', '--size_resize', nargs=2, type=int, default=[256, 256],
                        help='size of resizing input image')
    parser.add_argument('-sc', '--size_crop', nargs=2, type=int, default=[224, 224],
                        help='size of cropping after resize')
    parser.add_argument('-fh', '--flip_horz', action='store_true', help='flip horizontal')
    parser.add_argument('-fv', '--flip_vert', action='store_true', help='flip vertical')
    parser.add_argument('-tt', '--types_data', nargs='*', type=str, default=None)
    # feature extraction related
    parser.add_argument('-b', '--backbone', type=str,
                        default='torchvision.models.wide_resnet50_2',
                        help='specify torchvision model with the full path')
    parser.add_argument('-w', '--weight', type=str,
                        default='torchvision.models.Wide_ResNet50_2_Weights.IMAGENET1K_V1',
                        help='specify the trained weights of torchvision model with the full path')
    parser.add_argument('-lm', '--layer_map', nargs='+', type=str,
                        default=['layer2[-1]', 'layer3[-1]'],
                        help='specify layers to extract feature map')
    # patchification related
    parser.add_argument('-sp', '--size_patch', type=int, default=3,
                        help='patch pixel of feature map for increasing receptive field size')
    parser.add_argument('-de', '--dim_each_feat', type=int, default=1024,
                        help='dimension of extract feature (at 1st adaptive average pooling)')
    parser.add_argument('-dm', '--dim_merge_feat', type=int, default=1024,
                        help='dimension after layer feature merging (at 2nd adaptive average pooling)')
    # coreset related
    parser.add_argument('-s', '--seed', type=int, default=0,
                        help='specify a random-seed for k-center-greedy')
    parser.add_argument('-ns', '--num_split_seq', type=int, default=1,
                        help='percentage of coreset to all patch features')
    parser.add_argument('-pc', '--percentage_coreset', type=float, default=0.01,
                        help='percentage of coreset to all patch features')
    parser.add_argument('-ds', '--dim_sampling', type=int, default=128,
                        help='dimension to project features for sampling')
    parser.add_argument('-ni', '--num_initial_coreset', type=int, default=10,
                        help='number of samples to initially randomly select coreset')
    # Nearest-Neighbor related
    parser.add_argument('-k', '--k', type=int, default=5,
                        help='nearest neighbor\'s k for coreset searching')
    # post precessing related
    parser.add_argument('-pod', '--pixel_outer_decay', type=int, default=0,
                        help='number of outer pixels to decay anomaly score')

    args = parser.parse_args()
    return args


def apply_patchcore(type_data, feat_ext, patchcore, cfg_draw):
    print('\n----> PatchCore processing in %s start' % type_data)
    tic()

    # read images
    MVTecDataset(type_data)

    # reset neighbor
    patchcore.reset_neighbor()

    if args.verbose:
        # reset index
        idx_coreset_total = []

    # loop of split-sequential to apply k-center-greedy
    num_pitch = int(np.ceil(len(MVTecDataset.imgs_train) / patchcore.num_split_seq))
    for i_split in range(patchcore.num_split_seq):
        # extract features
        i_from = i_split * num_pitch
        i_to = min(((i_split + 1) * num_pitch), len(MVTecDataset.imgs_train))
        if (patchcore.num_split_seq > 1):
            print('[split%02d] image index range is %d~%d' % (i_split, i_from, (i_to - 1)))
        feat_train = feat_ext.extract(MVTecDataset.imgs_train[i_from:i_to])

        # coreset-reduced patch-feature memory bank
        idx_coreset = patchcore.compute_greedy_coreset_idx(feat_train)
        feat_train = feat_train[idx_coreset]

        # add feature as neighbor
        patchcore.add_neighbor(feat_train)

        if args.verbose:
            # stock index
            offset_split = i_from * feat_ext.HW_map()[0] * feat_ext.HW_map()[1]
            idx_coreset_total.append(idx_coreset + offset_split)

    if args.verbose:
        # concat index
        idx_coreset_total = np.hstack(idx_coreset_total)

    # extract features
    feat_test = {}
    for type_test in MVTecDataset.imgs_test.keys():
        feat_test[type_test] = feat_ext.extract(MVTecDataset.imgs_test[type_test],
                                                case='test (case:%s)' % type_test)

    # Sub-Image Anomaly Detection with Deep Pyramid Correspondences
    D, D_max, I = patchcore.localization(feat_test)

    # measure per image
    fpr_img, tpr_img, rocauc_img = calc_imagewise_metrics(D)
    print('%s imagewise ROCAUC: %.3f' % (type_data, rocauc_img))
    fpr_pix, tpr_pix, rocauc_pix = calc_pixelwise_metrics(D, MVTecDataset.gts_test)
    print('%s pixelwise ROCAUC: %.3f' % (type_data, rocauc_pix))

    toc(tag=('----> PatchCore processing in %s end, elapsed time' % type_data))

    draw_distance_graph(type_data, cfg_draw, D, rocauc_img)
    if args.verbose:
        draw_heatmap(type_data, cfg_draw, D, MVTecDataset.gts_test, D_max,
                     MVTecDataset.imgs_test, MVTecDataset.files_test,
                     idx_coreset_total, I, MVTecDataset.imgs_train, feat_ext.HW_map())

    return [fpr_img, tpr_img, rocauc_img, fpr_pix, tpr_pix, rocauc_pix]


def main(args):
    ConfigData(args)  # static define for speed-up
    cfg_feat = ConfigFeat(args)
    cfg_patchcore = ConfigPatchCore(args)
    cfg_draw = ConfigDraw(args)

    feat_ext = FeatExtract(cfg_feat)
    patchcore = PatchCore(cfg_patchcore, feat_ext.HW_map())

    os.makedirs(args.path_result, exist_ok=True)
    for type_data in ConfigData.types_data:
        os.makedirs(os.path.join(args.path_result, type_data), exist_ok=True)

    fpr_img = {}
    tpr_img = {}
    rocauc_img = {}
    fpr_pix = {}
    tpr_pix = {}
    rocauc_pix = {}

    # loop for types of data
    for type_data in ConfigData.types_data:
        result = apply_patchcore(type_data, feat_ext, patchcore, cfg_draw)

        fpr_img[type_data] = result[0]
        tpr_img[type_data] = result[1]
        rocauc_img[type_data] = result[2]

        fpr_pix[type_data] = result[3]
        tpr_pix[type_data] = result[4]
        rocauc_pix[type_data] = result[5]

    rocauc_img_mean = np.array([rocauc_img[type_data] for type_data in ConfigData.types_data])
    rocauc_img_mean = np.mean(rocauc_img_mean)
    rocauc_pix_mean = np.array([rocauc_pix[type_data] for type_data in ConfigData.types_data])
    rocauc_pix_mean = np.mean(rocauc_pix_mean)

    draw_roc_curve(cfg_draw, fpr_img, tpr_img, rocauc_img, rocauc_img_mean,
                             fpr_pix, tpr_pix, rocauc_pix, rocauc_pix_mean)

    for type_data in ConfigData.types_data:
        print('rocauc_img[%s] = %.3f' % (type_data, rocauc_img[type_data]))
    print('rocauc_img[mean] = %.3f' % rocauc_img_mean)
    for type_data in ConfigData.types_data:
        print('rocauc_pix[%s] = %.3f' % (type_data, rocauc_pix[type_data]))
    print('rocauc_pix[mean] = %.3f' % rocauc_pix_mean)


if __name__ == '__main__':
    args = arg_parser()
    main(args)

main.py
any-tech/PatchCore-ex
import os
import numpy as np
from tqdm import tqdm
import torch
import torch.nn.functional as F
import torchvision
from torchinfo import summary


class FeatExtract:
    def __init__(self, cfg_feat):
        self.device = cfg_feat.device
        self.batch_size = cfg_feat.batch_size
        self.shape_input = cfg_feat.SHAPE_INPUT
        self.layer_map = cfg_feat.layer_map
        self.size_patch = cfg_feat.size_patch
        self.dim_each_feat = cfg_feat.dim_each_feat
        self.dim_merge_feat = cfg_feat.dim_merge_feat
        self.MEAN = cfg_feat.MEAN
        self.STD = cfg_feat.STD

        self.padding = int((self.size_patch - 1) / 2)
        self.stride = 1  # fixed temporarily...

        code = 'self.backbone = %s(weights=%s)' % (cfg_feat.backbone, cfg_feat.weight)
        exec(code)
        self.backbone.eval()
        self.backbone.to(self.device)
        summary(self.backbone, input_size=(1, 3, *self.shape_input))

        self.feat = []
        for layer_map in self.layer_map:
            code = 'self.backbone.%s.register_forward_hook(self.hook)' % layer_map
            exec(code)

        # dummy forward
        x = torch.zeros(1, 3, self.shape_input[0], self.shape_input[1])  # RGB
        x = x.to(self.device)
        self.feat = []
        with torch.no_grad():
            _ = self.backbone(x)

        # https://github.com/amazon-science/patchcore-inspection/blob/main/src/patchcore/patchcore.py#L295
        self.unfolder = torch.nn.Unfold(kernel_size=self.size_patch, stride=self.stride,
                                        padding=self.padding, dilation=1)
        self.patch_shapes = []
        for i in range(len(self.feat)):
            number_of_total_patches = []
            for s in self.feat[i].shape[-2:]:
                n_patches = (
                    s + 2 * self.padding - 1 * (self.size_patch - 1) - 1
                ) / self.stride + 1
                number_of_total_patches.append(int(n_patches))
            self.patch_shapes.append(number_of_total_patches)

    def hook(self, module, input, output):
        self.feat.append(output.detach().cpu())

    def HW_map(self):
        return self.patch_shapes[0]

    def normalize(self, input):
        x = torch.from_numpy(input.astype(np.float32))
        x = x.to(self.device)
        x = x / 255
        x = x - self.MEAN
        x = x / self.STD
        x = x.unsqueeze(0).permute(0, 3, 1, 2)
        return x

    # return : is_train=True->torch.Tensor, is_train=False->dict
    def extract(self, imgs, case='train (case:good)', batch_size_patchfy=50):
        # feature extract for train and aggregate split-image for explain
        x_batch = []
        self.feat = []
        for i_img in tqdm(range(len(imgs)),
                          desc='extract feature for %s' % case):
            img = imgs[i_img]
            x = self.normalize(img)
            x_batch.append(x)

            if ((len(x_batch) == self.batch_size) | (i_img == (len(imgs) - 1))):
                with torch.no_grad():
                    _ = self.backbone(torch.vstack(x_batch))
                x_batch = []

        # adjust
        feat = []
        num_layer = len(self.layer_map)
        for i_layer_map in range(num_layer):
            feat.append(torch.vstack(self.feat[i_layer_map::num_layer]))

        # patchfy (consider out of memory)
        num_patch_per_image = self.HW_map()[0] * self.HW_map()[1]
        num_patch = len(imgs) * num_patch_per_image
        feat_patchfy = torch.zeros(num_patch, self.dim_merge_feat)

        num_patchfy_process = (len(feat) * 3) + 1
        num_iter = np.ceil(len(imgs) / batch_size_patchfy)
        pbar = tqdm(total=int(num_patchfy_process * num_iter), desc='patchfy feature')
        for i_batch in range(0, len(imgs), batch_size_patchfy):
            feat_tmp = []
            for feat_layer in feat:
                feat_tmp.append(feat_layer[i_batch:(i_batch + batch_size_patchfy)])
            i_from = i_batch * num_patch_per_image
            i_to = (i_batch + batch_size_patchfy) * num_patch_per_image
            feat_patchfy[i_from:i_to] = self.patchfy(feat_tmp, pbar)
        pbar.close()

        return feat_patchfy

    def patchfy(self, feat, pbar, batch_size_interp=2000):

        with torch.no_grad():
            # unfold
            for i in range(len(feat)):
                _feat = feat[i]
                BC_before_unfold = _feat.shape[:2]
                # (B, C, H, W) -> (B, CPHPW, HW)
                _feat = self.unfolder(_feat)
                # (B, CPHPW, HW) -> (B, C, PH, PW, HW)
                _feat = _feat.reshape(*BC_before_unfold,
                                      self.size_patch, self.size_patch, -1)
                # (B, C, PH, PW, HW) -> (B, HW, C, PW, HW)
                _feat = _feat.permute(0, 4, 1, 2, 3)
                feat[i] = _feat
                pbar.update(1)

            # expand small feat to fit large features
            for i in range(1, len(feat)):
                _feat = feat[i]
                patch_dims = self.patch_shapes[i]
                # (B, HW, C, PW, HW) -> (B, H, W, C, PH, PW)
                _feat = _feat.reshape(_feat.shape[0], patch_dims[0],
                                      patch_dims[1], *_feat.shape[2:])
                # (B, H, W, C, PH, PW) -> (B, C, PH, PW, H, W)
                _feat = _feat.permute(0, -3, -2, -1, 1, 2)
                perm_base_shape = _feat.shape
                # (B, C, PH, PW, H, W) -> (BCPHPW, H, W)
                _feat = _feat.reshape(-1, *_feat.shape[-2:])
                # (BCPHPW, H, W) -> (BCPHPW, H_max, W_max)
                __feat = torch.zeros([len(_feat), self.HW_map()[0], self.HW_map()[1]])
                for i_batch in range(0, len(_feat), batch_size_interp):
                    feat_tmp = _feat[i_batch:(i_batch + batch_size_interp)]
                    feat_tmp = feat_tmp.unsqueeze(1).to(self.device)
                    feat_tmp = F.interpolate(feat_tmp,
                                             size=(self.HW_map()[0], self.HW_map()[1]),
                                             mode="bilinear", align_corners=False)
                    __feat[i_batch:(i_batch + batch_size_interp)] = feat_tmp.squeeze(1).cpu()
                _feat = __feat
                # _feat = F.interpolate(_feat.unsqueeze(1),
                #                       size=(self.HW_map()[0], self.HW_map()[1]),
                #                       mode="bilinear", align_corners=False)
                # _feat = _feat.squeeze(1)
                # for i_batch in range(0, len(_feat), 10000):
                #     print(torch.sum(torch.abs(_feat[i_batch:(i_batch + 10000)] - __feat[i_batch:(i_batch + 10000)])))
                # (BCPHPW, H_max, W_max) -> (B, C, PH, PW, H_max, W_max)
                _feat = _feat.reshape(*perm_base_shape[:-2],
                                      self.HW_map()[0], self.HW_map()[1])
                # (B, C, PH, PW, H_max, W_max) -> (B, H_max, W_max, C, PH, PW)
                _feat = _feat.permute(0, -2, -1, 1, 2, 3)
                # (B, H_max, W_max, C, PH, PW) -> (B, H_maxW_max, C, PH, PW)
                _feat = _feat.reshape(len(_feat), -1, *_feat.shape[-3:])
                feat[i] = _feat
                pbar.update(1)

            # aggregate feature vectors
            # (B, H, W, C, PH, PW) -> (BHW, C, PH, PW)
            feat = [x.reshape(-1, *x.shape[-3:]) for x in feat]
            pbar.update(1)

            # adaptive average pooling for each feature vector
            for i in range(len(feat)):
                _feat = feat[i]
                # (BHW, C, PH, PW) -> (BHW, 1, CPHPW)
                _feat = _feat.reshape(len(_feat), 1, -1)
                # (BHW, 1, CPHPW) -> (BHW, D_e)
                _feat = F.adaptive_avg_pool1d(_feat,
                                              self.dim_each_feat).squeeze(1)
                feat[i] = _feat
                pbar.update(1)

            # concat the two feature vectors and adaptive average pooling
            # (BHW, D_e) -> (BHW, D_e*2)
            feat = torch.stack(feat, dim=1)
            """Returns reshaped and average pooled feat."""
            # batchsize x number_of_layers x input_dim -> batchsize x target_dim
            # (BHW, D_e*2) -> (BHW, D_m)
            feat = feat.reshape(len(feat), 1, -1)
            feat = F.adaptive_avg_pool1d(feat, self.dim_merge_feat)
            feat = feat.reshape(len(feat), -1)
            pbar.update(1)

        return feat

models/feat_extract.py
any-tech/PatchCore-ex
import numpy as np
import cv2
from tqdm import tqdm
import torch
import faiss
from scipy.ndimage import gaussian_filter


class PatchCore:
    def __init__(self, cfg_patchcore, HW_map):
        self.device = cfg_patchcore.device
        self.k = cfg_patchcore.k
        self.dim_coreset_feat = cfg_patchcore.dim_coreset_feat
        self.num_split_seq = cfg_patchcore.num_split_seq
        self.percentage_coreset = cfg_patchcore.percentage_coreset
        self.dim_sampling = cfg_patchcore.dim_sampling
        self.num_initial_coreset = cfg_patchcore.num_initial_coreset
        self.seed = cfg_patchcore.seed
        self.shape_stretch = cfg_patchcore.shape_stretch
        self.pixel_outer_decay = cfg_patchcore.pixel_outer_decay
        self.HW_map = HW_map

        # prep knn index
        self.index_feat = faiss.GpuIndexFlatL2(faiss.StandardGpuResources(),
                                               self.dim_coreset_feat,
                                               faiss.GpuIndexFlatConfig())

        # prep mapper
        self.mapper = torch.nn.Linear(self.dim_coreset_feat, self.dim_sampling,
                                      bias=False).to(self.device)

    def compute_greedy_coreset_idx(self, feat):
        feat = feat.to(self.device)
        with torch.no_grad():
            feat_proj = self.mapper(feat)

        _num_initial_coreset = np.clip(self.num_initial_coreset,
                                       None, len(feat_proj))
        np.random.seed(self.seed)
        start_points = np.random.choice(len(feat_proj), _num_initial_coreset,
                                        replace=False).tolist()

        # computes batchwise Euclidean distances using PyTorch
        mat_A = feat_proj
        mat_B = feat_proj[start_points]
        A_x_A = mat_A.unsqueeze(1).bmm(mat_A.unsqueeze(2)).reshape(-1, 1)
        B_x_B = mat_B.unsqueeze(1).bmm(mat_B.unsqueeze(2)).reshape(1, -1)
        A_x_B = mat_A.mm(mat_B.T)
        # not need sqrt
        mat_dist = (-2 * A_x_B + A_x_A + B_x_B).clamp(0, None)

        dist_coreset_anchor = torch.mean(mat_dist, axis=-1, keepdims=True)

        idx_coreset = []
        num_coreset_samples = int(len(feat_proj) * self.percentage_coreset)

        with torch.no_grad():
            for _ in tqdm(range(num_coreset_samples), desc="sampling"):
                idx_select = torch.argmax(dist_coreset_anchor).item()
                idx_coreset.append(idx_select)

                mat_A = feat_proj
                mat_B = feat_proj[[idx_select]]
                # computes batchwise Euclidean distances using PyTorch
                A_x_A = mat_A.unsqueeze(1).bmm(mat_A.unsqueeze(2)).reshape(-1, 1)
                B_x_B = mat_B.unsqueeze(1).bmm(mat_B.unsqueeze(2)).reshape(1, -1)
                A_x_B = mat_A.mm(mat_B.T)
                # not need sqrt
                mat_select_dist = (-2 * A_x_B + A_x_A + B_x_B).clamp(0, None)

                dist_coreset_anchor = torch.cat([dist_coreset_anchor,
                                                 mat_select_dist], dim=-1)
                dist_coreset_anchor = torch.min(dist_coreset_anchor,
                                                dim=1).values.reshape(-1, 1)

        idx_coreset = np.array(idx_coreset)
        return idx_coreset

    def reset_neighbor(self):
        self.index_feat.reset()

    def add_neighbor(self, feat_train):
        self.index_feat.add(feat_train.numpy())

    def localization(self, feat_test):
        D = {}
        D_max = -9999
        I = {}
        # loop for test cases
        for type_test in feat_test.keys():
            D[type_test] = []
            I[type_test] = []

            # loop for test data
            _feat_test = feat_test[type_test]
            _feat_test = _feat_test.reshape(-1, (self.HW_map[0] * self.HW_map[1]),
                                            self.dim_coreset_feat)
            _feat_test = _feat_test.numpy()
            num_data = len(_feat_test)
            for i in tqdm(range(num_data),
                          desc='localization (case:%s)' % type_test):
                # measure distance pixelwise
                score_map, I_tmp = self.measure_dist_pixelwise(feat_test=_feat_test[i])
                # adjust score of outer-pixel (provisional heuristic algorithm)
                if (self.pixel_outer_decay > 0):
                    score_map[:self.pixel_outer_decay, :] *= 0.6
                    score_map[-self.pixel_outer_decay:, :] *= 0.6
                    score_map[:, :self.pixel_outer_decay] *= 0.6
                    score_map[:, -self.pixel_outer_decay:] *= 0.6
                # stock score map
                D[type_test].append(score_map)
                D_max = max(D_max, np.max(score_map))
                I[type_test].append(I_tmp)

            # cast list to numpy array
            D[type_test] = np.array(D[type_test])
            I[type_test] = np.array(I[type_test])

        return D, D_max, I

    def measure_dist_pixelwise(self, feat_test):
        # k nearest neighbor
        D, I = self.index_feat.search(feat_test, self.k)
        D = np.mean(D, axis=-1)

        # transform to scoremap
        score_map = D.reshape(*self.HW_map)
        score_map = cv2.resize(score_map, (self.shape_stretch[1],
                                           self.shape_stretch[0]))

        # apply gaussian smoothing on the score map
        score_map_smooth = gaussian_filter(score_map, sigma=4)

        return score_map_smooth, I

models/patchcore.py
any-tech/PatchCore-ex
import os
import numpy as np
import torch

# https://pytorch.org/vision/main/models/generated/torchvision.models.wide_resnet50_2.html
MEAN = torch.FloatTensor([[[0.485, 0.456, 0.406]]])
STD = torch.FloatTensor([[[0.229, 0.224, 0.225]]])


class ConfigData:
    @classmethod
    def __init__(cls, args):
        # file reading related
        cls.path_parent = args.path_parent
        assert os.path.exists(cls.path_parent)
        cls.num_cpu_max = args.num_cpu_max
        cls.shuffle = (args.num_split_seq > 1)  # for k-center-greedy split
        cls.seed = args.seed

        # input format related
        cls.SHAPE_MIDDLE = (args.size_resize[0], args.size_resize[1])  # (H, W)
        cls.SHAPE_INPUT = (args.size_crop[0], args.size_crop[1])  # (H, W)
        cls.pixel_cut = (int((cls.SHAPE_MIDDLE[0] - cls.SHAPE_INPUT[0]) / 2),
                         int((cls.SHAPE_MIDDLE[1] - cls.SHAPE_INPUT[1]) / 2))  # (H, W)

        # augmantation related
        cls.flip_horz = args.flip_horz
        cls.flip_vert = args.flip_vert

        # collect types of data
        if args.types_data is None:
            types_data = [d for d in os.listdir(args.path_parent)
                          if os.path.isdir(os.path.join(args.path_parent, d))]
            cls.types_data = np.sort(np.array(types_data))
        else:
            for type_data in args.types_data:
                assert os.path.exists(os.path.join(args.path_parent, type_data))
            cls.types_data = np.sort(np.array(args.types_data))


class ConfigFeat:
    def __init__(self, args):
        # adjsut to environment
        if args.cpu:
            self.device = torch.device('cpu')
        else:
            self.device = torch.device('cuda:0')

        # batch-size for feature extraction by ImageNet model
        self.batch_size = args.batch_size

        # input format related
        self.SHAPE_INPUT = (args.size_crop[0], args.size_crop[1])  # (H, W)

        # base network
        self.backbone = args.backbone
        self.weight = args.weight

        # layer specification
        self.layer_map = args.layer_map

        # patch pixel of feature map for increasing receptive field size
        self.size_patch = args.size_patch

        # dimension of each layer feature (at 1st adaptive average pooling)
        self.dim_each_feat = args.dim_each_feat
        # dimension after layer feature merging (at 2nd adaptive average pooling)
        self.dim_merge_feat = args.dim_merge_feat

        # adjust to the network's learning policy and the data conditions
        self.MEAN = MEAN.to(self.device)
        self.STD = STD.to(self.device)


class ConfigPatchCore:
    def __init__(self, args):
        # adjsut to environment
        if args.cpu:
            self.device = torch.device('cpu')
        else:
            self.device = torch.device('cuda:0')

        # dimension after layer feature merging (at 2nd adaptive average pooling)
        self.dim_coreset_feat = args.dim_merge_feat

        # number split-sequential to apply k-center-greedy
        self.num_split_seq = args.num_split_seq
        # percentage of coreset to all patch features
        self.percentage_coreset = args.percentage_coreset
        # dimension to project features for sampling
        self.dim_sampling = args.dim_sampling
        # number of samples to initially randomly select coreset
        self.num_initial_coreset = args.num_initial_coreset
        # random-seed for k-center-greedy
        self.seed = args.seed

        # number of nearest neighbor to get patch images
        self.k = args.k

        # input format related
        self.shape_stretch = (args.size_crop[0], args.size_crop[1])  # (H, W)

        # consideration for the outer edge
        self.pixel_outer_decay = args.pixel_outer_decay


class ConfigDraw:
    def __init__(self, args):
        # output detail or not (take a long time...)
        self.verbose = args.verbose

        # output filename related
        self.k = args.k
        self.percentage_coreset = args.percentage_coreset

        # receptive field size
        self.size_receptive_field = args.size_receptive_field

        # aspect_ratio of output figure
        self.aspect_figure = args.size_crop[1] / args.size_crop[0]  # W / H
        self.aspect_figure = np.round(self.aspect_figure, decimals=1)

        # visualize mode
        self.mode_visualize = args.mode_visualize

        # output path of figure
        self.path_result = args.path_result

utils/config.py
any-tech/PatchCore-ex
import numpy as np
import numba
from tqdm import tqdm
from sklearn.metrics import roc_curve
# from sklearn.metrics import roc_auc_score


def calc_imagewise_metrics(D, type_normal='good'):
    D_list = []
    y_list = []

    pbar = tqdm(total=int(len(D.keys()) + 3),
                desc='calculate imagewise metrics')
    for type_test in D.keys():
        for i in range(len(D[type_test])):
            D_tmp = np.max(D[type_test][i])
            y_tmp = int(type_test != type_normal)

            D_list.append(D_tmp)
            y_list.append(y_tmp)
        pbar.update(1)

    D_flatten_list = np.array(D_list).reshape(-1)
    y_flatten_list = np.array(y_list).reshape(-1)
    pbar.update(1)

    fpr, tpr, _ = roc_curve(y_flatten_list, D_flatten_list)
    pbar.update(1)

    rocauc = roc_auc_score(y_flatten_list, D_flatten_list)
    pbar.update(1)
    pbar.close()

    return fpr, tpr, rocauc


def calc_pixelwise_metrics(D, y):
    D_list = []
    y_list = []

    pbar = tqdm(total=int(len(D.keys()) + 3),
                desc='calculate pixelwise metrics')
    for type_test in D.keys():
        for i in range(len(D[type_test])):
            D_tmp = D[type_test][i]
            y_tmp = y[type_test][i]

            D_list.append(D_tmp)
            y_list.append(y_tmp)
        pbar.update(1)

    D_flatten_list = np.array(D_list).reshape(-1)
    pbar.update(1)

    y_flatten_list = np.array(y_list).reshape(-1)
    pbar.update(1)

    fpr, tpr, _ = roc_curve(y_flatten_list, D_flatten_list)
    rocauc = roc_auc_score(y_flatten_list, D_flatten_list)
    pbar.update(1)
    pbar.close()

    return fpr, tpr, rocauc


# https://github.com/diditforlulz273/fastauc/blob/main/fastauc/fast_auc.py
def roc_auc_score(y_true, y_score):
    # binary clf curve
    y_true = (y_true == 1)

    desc_score_indices = np.argsort(y_score, kind="mergesort")[::-1]
    y_score = y_score[desc_score_indices]
    y_true = y_true[desc_score_indices]

    distinct_value_indices = np.where(np.diff(y_score))[0]
    threshold_idxs = np.r_[distinct_value_indices, y_true.size - 1]

    tps = np.cumsum(y_true)[threshold_idxs]
    fps = 1 + threshold_idxs - tps

    # roc
    tps = np.r_[0, tps]
    fps = np.r_[0, fps]

    if fps[-1] <= 0 or tps[-1] <= 0:
        return np.nan

    # auc
    direction = 1
    dx = np.diff(fps)
    if np.any(dx < 0):
        if np.all(dx <= 0):
            direction = -1
        else:
            return 'error'

    area = direction * np.trapz(tps, fps) / (tps[-1] * fps[-1])

    return area

utils/metrics.py
any-tech/PatchCore-ex
import time


def tic():
    # require to import time
    global start_time_tictoc
    start_time_tictoc = time.time()


def toc(tag="elapsed time"):
    if "start_time_tictoc" in globals():
        print("{}: {:.1f} [sec]".format(tag, time.time() - start_time_tictoc))
    else:
        print("tic has not been called")

utils/tictoc.py
any-tech/PatchCore-ex
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from tqdm import tqdm


# https://github.com/gsurma/cnn_explainer/blob/main/utils.py
def overlay_heatmap_on_image(img, heatmap, ratio_img=0.5):
    img = img.astype(np.float32)

    heatmap = 1 - np.clip(heatmap, 0, 1)
    heatmap = cv2.applyColorMap((heatmap * 255).astype(np.uint8), cv2.COLORMAP_JET)
    heatmap = heatmap.astype(np.float32)

    overlay = (img * ratio_img) + (heatmap * (1 - ratio_img))
    overlay = np.clip(overlay, 0, 255)
    overlay = overlay.astype(np.uint8)
    return overlay


def draw_distance_graph(type_data, cfg_draw, D, rocauc_img):
    D_list = {}
    for type_test in D.keys():
        D_list[type_test] = []
        for i in range(len(D[type_test])):
            D_tmp = np.max(D[type_test][i])
            D_list[type_test].append(D_tmp)
        D_list[type_test] = np.array(D_list[type_test])

    plt.figure(figsize=(10, 8), dpi=100, facecolor='white')

    # 'good' 1st
    N_test = 0
    type_test = 'good'
    plt.subplot(2, 1, 1)
    plt.scatter((np.arange(len(D_list[type_test])) + N_test), D_list[type_test],
                alpha=0.5, label=type_test)
    plt.subplot(2, 1, 2)
    plt.hist(D_list[type_test], alpha=0.5, label=type_test, bins=10)

    # other than 'good'
    N_test += len(D_list[type_test])
    types_test = np.array([k for k in D_list.keys() if k != 'good'])
    for type_test in types_test:
        plt.subplot(2, 1, 1)
        plt.scatter((np.arange(len(D_list[type_test])) + N_test), D_list[type_test],
                    alpha=0.5, label=type_test)
        plt.subplot(2, 1, 2)
        plt.hist(D_list[type_test], alpha=0.5, label=type_test, bins=10)
        N_test += len(D_list[type_test])

    plt.subplot(2, 1, 1)
    plt.title('imagewise anomaly detection accuracy (ROCAUC %%) : %.3f' % rocauc_img)
    plt.grid()
    plt.legend()
    plt.subplot(2, 1, 2)
    plt.grid()
    plt.legend()
    plt.gcf().tight_layout()
    plt.gcf().savefig(os.path.join(cfg_draw.path_result, type_data,
                                   ('pred-dist_%s_p%04d_k%02d_r%04d.png' %
                                    (type_data, (cfg_draw.percentage_coreset * 1000),
                                     cfg_draw.k, round(rocauc_img * 1000)))))
    plt.clf()
    plt.close()


def draw_heatmap(type_data, cfg_draw, D, y, D_max, imgs_test, files_test,
                 idx_coreset, I, imgs_train, HW_map):
    for type_test in D.keys():
        for i in tqdm(range(len(D[type_test])),
                      desc='[verbose mode] visualize localization (case:%s)' % type_test):
            file = files_test[type_test][i]
            img = imgs_test[type_test][i]
            score_map = D[type_test][i]
            idx_patch = idx_coreset[I[type_test][i][:, 0]]
            score_max = D_max
            gt = y[type_test][i]
            img_patch = pickup_patch(idx_patch, imgs_train, HW_map,
                                     cfg_draw.size_receptive_field)

            plt.figure(figsize=(10 * max(1, cfg_draw.aspect_figure), 18),
                       dpi=100, facecolor='white')
            plt.rcParams['font.size'] = 10

            if (cfg_draw.mode_visualize == 'eval'):
                plt.subplot2grid((7, 3), (0, 0), rowspan=1, colspan=1)
                plt.imshow(img)
                plt.title('%s : %s' % (file.split('/')[-2], file.split('/')[-1]))
                plt.subplot2grid((7, 3), (0, 1), rowspan=1, colspan=1)
                plt.imshow(gt)
                plt.subplot2grid((7, 3), (0, 2), rowspan=1, colspan=1)
                plt.imshow(score_map)
                plt.colorbar()
                plt.title('max score : %.2f' % score_max)
                plt.subplot2grid((42, 2), (7, 0), rowspan=10, colspan=1)
                plt.imshow(overlay_heatmap_on_image(img, (score_map / score_max)))
                plt.subplot2grid((42, 2), (7, 1), rowspan=10, colspan=1)
                plt.imshow((img.astype(np.float32) *
                            (score_map / score_max)[..., None]).astype(np.uint8))

                plt.subplot2grid((21, 1), (10, 0), rowspan=11, colspan=1)
                plt.imshow(img_patch, interpolation='none')
                plt.title('patch images created with top1-NN')
            elif (cfg_draw.mode_visualize == 'infer'):
                plt.subplot2grid((5, 2), (0, 0), rowspan=1, colspan=1)
                plt.imshow(img)
                plt.title('%s : %s' % (file.split('/')[-2], file.split('/')[-1]))
                plt.subplot2grid((5, 2), (0, 1), rowspan=1, colspan=1)
                plt.imshow(score_map)
                plt.colorbar()
                plt.title('max score : %.2f' % score_max)
                plt.subplot2grid((5, 1), (1, 0), rowspan=2, colspan=1)
                plt.imshow(overlay_heatmap_on_image(img, (score_map / score_max)))

                plt.subplot2grid((5, 1), (3, 0), rowspan=2, colspan=1)
                plt.imshow(img_patch, interpolation='none')
                plt.title('patch images created with top1-NN')

            score_tmp = np.max(score_map) / score_max * 100
            filename_out = os.path.join(cfg_draw.path_result, type_data,
                                        ('localization_%s_%s_%s_p%04d_k%02d_s%03d.png' %
                                         (type_data, type_test, 
                                          os.path.basename(file).split('.')[0],
                                          (cfg_draw.percentage_coreset * 1000),
                                          cfg_draw.k, round(score_tmp))))

            plt.gcf().savefig(filename_out)
            plt.clf()
            plt.close()


def pickup_patch(idx_patch, imgs, HW_map, size_receptive_field):
    # get input shape
    h = imgs.shape[-3]
    w = imgs.shape[-2]
    # calculate half size for split
    h_half = int((size_receptive_field - 1) / 2)
    w_half = int((size_receptive_field - 1) / 2)
    # calculate center-coordinates of split-image
    y_pitch = np.arange(0, (h - 1 + 1e-10), ((h - 1) / (HW_map[0] - 1)))
    y_pitch = np.round(y_pitch).astype(np.int16)
    y_pitch = y_pitch + h_half
    x_pitch = np.arange(0, (w - 1 + 1e-10), ((w - 1) / (HW_map[1] - 1)))
    x_pitch = np.round(x_pitch).astype(np.int16)
    x_pitch = x_pitch + w_half
    # padding to normal images
    imgs = np.pad(imgs, ((0, 0), (h_half, h_half), (w_half, w_half), (0, 0)))

    # build blank image for output
    img_patch = np.zeros([(HW_map[0] * size_receptive_field),
                          (HW_map[1] * size_receptive_field), 3], dtype=np.uint8)
    i_y = 0
    i_x = 0

    # loop of patch feature index
    for i_patch in idx_patch:
        i_img = i_patch // (HW_map[0] * HW_map[1])
        i_HW = i_patch % (HW_map[0] * HW_map[1])
        i_H = i_HW // HW_map[1]
        i_W = i_HW % HW_map[1]

        img = imgs[i_img]
        y = y_pitch[i_H]
        x = x_pitch[i_W]
        img_piece = img[(y - h_half):(y + h_half + 1),
                        (x - w_half):(x + w_half + 1)]

        y = i_y * size_receptive_field
        x = i_x * size_receptive_field
        img_patch[y:(y + size_receptive_field),
                  x:(x + size_receptive_field)] = img_piece
        i_x += 1
        if (i_x >= HW_map[1]):
            i_x = 0
            i_y += 1

    return img_patch


def draw_roc_curve(cfg_draw, fpr_img, tpr_img, rocauc_img, rocauc_img_mean,
                             fpr_pix, tpr_pix, rocauc_pix, rocauc_pix_mean):
    plt.figure(figsize=(12, 6), dpi=100, facecolor='white')
    for type_data in fpr_img.keys():
        plt.subplot(1, 2, 1)
        plt.plot(fpr_img[type_data], tpr_img[type_data],
                 label='%s ROCAUC: %.3f' % (type_data, rocauc_img[type_data]))
        plt.subplot(1, 2, 2)
        plt.plot(fpr_pix[type_data], tpr_pix[type_data],
                 label='%s ROCAUC: %.3f' % (type_data, rocauc_pix[type_data]))

    plt.subplot(1, 2, 1)
    plt.title('imagewise anomaly detection accuracy (ROCAUC %%) : mean %.3f' % rocauc_img_mean)
    plt.grid()
    plt.legend()
    plt.subplot(1, 2, 2)
    plt.title('pixelwise anomaly detection accuracy (ROCAUC %%) : mean %.3f' % rocauc_pix_mean)
    plt.grid()
    plt.legend()
    plt.gcf().tight_layout()
    plt.gcf().savefig(os.path.join(cfg_draw.path_result,
                                   ('roc-curve_p%04d_k%02d_rim%04d_rpm%04d.png' %
                                    ((cfg_draw.percentage_coreset * 1000), cfg_draw.k,
                                     round(rocauc_img_mean * 1000),
                                     round(rocauc_pix_mean * 1000)))))
    plt.clf()
    plt.close()

utils/visualize.py
any-tech/PatchCore-ex