SAPADE_main.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.4"
    },
    "colab": {
      "name": "SAPADE main.ipynb",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/danielsoy/data/blob/main/SAPADE_main.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0kJ2ySjPTehE"
      },
      "source": [
        "import argparse\n",
        "import numpy as np\n",
        "import os\n",
        "import pickle\n",
        "from tqdm import tqdm\n",
        "from collections import OrderedDict\n",
        "from sklearn.metrics import roc_auc_score\n",
        "from sklearn.metrics import roc_curve\n",
        "from sklearn.metrics import precision_recall_curve\n",
        "from scipy.ndimage import gaussian_filter\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cyirr_J2TehG"
      },
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.models import wide_resnet50_2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UdjcjSphTehH"
      },
      "source": [
        "import datasets.mvtec as mvtec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5qmM1LTHTehI"
      },
      "source": [
        "def parse_args():\n",
        "    parser = argparse.ArgumentParser('SPADE')\n",
        "    parser.add_argument(\"--top_k\", type=int, default=5)\n",
        "    parser.add_argument(\"--save_path\", type=str, default=\"./result\")\n",
        "    return parser.parse_args()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qJBd3E-iTehJ"
      },
      "source": [
        "def main():\n",
        "    args = parse_args()\n",
        "\n",
        "    # device setup\n",
        "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "\n",
        "    # load model\n",
        "    model = wide_resnet50_2(pretrained=True, progress=True)\n",
        "    model.to(device)\n",
        "    model.eval()\n",
        "\n",
        "    # set model's intermediate outputs\n",
        "    outputs = []\n",
        "    def hook(module, input, output):\n",
        "        outputs.append(output)\n",
        "    model.layer1[-1].register_forward_hook(hook)\n",
        "    model.layer2[-1].register_forward_hook(hook)\n",
        "    model.layer3[-1].register_forward_hook(hook)\n",
        "    model.avgpool.register_forward_hook(hook)\n",
        "    os.makedirs(os.path.join(args.save_path, 'temp'), exist_ok=True)\n",
        "    fig, ax = plt.subplots(1, 2, figsize=(20, 10))\n",
        "    fig_img_rocauc = ax[0]\n",
        "    fig_pixel_rocauc = ax[1]\n",
        "    total_roc_auc = []\n",
        "    total_pixel_roc_auc = []\n",
        "    for class_name in mvtec.CLASS_NAMES:\n",
        "        train_dataset = mvtec.MVTecDataset(class_name=class_name, is_train=True)\n",
        "        train_dataloader = DataLoader(train_dataset, batch_size=32, pin_memory=True)\n",
        "        test_dataset = mvtec.MVTecDataset(class_name=class_name, is_train=False)\n",
        "        test_dataloader = DataLoader(test_dataset, batch_size=32, pin_memory=True)\n",
        "        train_outputs = OrderedDict([('layer1', []), ('layer2', []), ('layer3', []), ('avgpool', [])])\n",
        "        test_outputs = OrderedDict([('layer1', []), ('layer2', []), ('layer3', []), ('avgpool', [])])\n",
        "\n",
        "        # extract train set features\n",
        "        train_feature_filepath = os.path.join(args.save_path, 'temp', 'train_%s.pkl' % class_name)\n",
        "        if not os.path.exists(train_feature_filepath):\n",
        "            for (x, y, mask) in tqdm(train_dataloader, '| feature extraction | train | %s |' % class_name):\n",
        "                # model prediction\n",
        "                with torch.no_grad():\n",
        "                    pred = model(x.to(device))\n",
        "                # get intermediate layer outputs\n",
        "                for k, v in zip(train_outputs.keys(), outputs):\n",
        "                    train_outputs[k].append(v)\n",
        "                # initialize hook outputs\n",
        "                outputs = []\n",
        "            for k, v in train_outputs.items():\n",
        "                train_outputs[k] = torch.cat(v, 0)\n",
        "            # save extracted feature\n",
        "            with open(train_feature_filepath, 'wb') as f:\n",
        "                pickle.dump(train_outputs, f)\n",
        "        else:\n",
        "            print('load train set feature from: %s' % train_feature_filepath)\n",
        "            with open(train_feature_filepath, 'rb') as f:\n",
        "                train_outputs = pickle.load(f)\n",
        "        gt_list = []\n",
        "        gt_mask_list = []\n",
        "        test_imgs = []\n",
        "\n",
        "        # extract test set features\n",
        "        for (x, y, mask) in tqdm(test_dataloader, '| feature extraction | test | %s |' % class_name):\n",
        "            test_imgs.extend(x.cpu().detach().numpy())\n",
        "            gt_list.extend(y.cpu().detach().numpy())\n",
        "            gt_mask_list.extend(mask.cpu().detach().numpy())\n",
        "            # model prediction\n",
        "            with torch.no_grad():\n",
        "                pred = model(x.to(device))\n",
        "            # get intermediate layer outputs\n",
        "            for k, v in zip(test_outputs.keys(), outputs):\n",
        "                test_outputs[k].append(v)\n",
        "            # initialize hook outputs\n",
        "            outputs = []\n",
        "        for k, v in test_outputs.items():\n",
        "            test_outputs[k] = torch.cat(v, 0)\n",
        "\n",
        "        # calculate distance matrix\n",
        "        dist_matrix = calc_dist_matrix(torch.flatten(test_outputs['avgpool'], 1),\n",
        "                                       torch.flatten(train_outputs['avgpool'], 1))\n",
        "\n",
        "        # select K nearest neighbor and take average\n",
        "        topk_values, topk_indexes = torch.topk(dist_matrix, k=args.top_k, dim=1, largest=False)\n",
        "        scores = torch.mean(topk_values, 1).cpu().detach().numpy()\n",
        "\n",
        "        # calculate image-level ROC AUC score\n",
        "        fpr, tpr, _ = roc_curve(gt_list, scores)\n",
        "        roc_auc = roc_auc_score(gt_list, scores)\n",
        "        total_roc_auc.append(roc_auc)\n",
        "        print('%s ROCAUC: %.3f' % (class_name, roc_auc))\n",
        "        fig_img_rocauc.plot(fpr, tpr, label='%s ROCAUC: %.3f' % (class_name, roc_auc))\n",
        "        score_map_list = []\n",
        "        for t_idx in tqdm(range(test_outputs['avgpool'].shape[0]), '| localization | test | %s |' % class_name):\n",
        "            score_maps = []\n",
        "            for layer_name in ['layer1', 'layer2', 'layer3']:  # for each layer\n",
        "                # construct a gallery of features at all pixel locations of the K nearest neighbors\n",
        "                topk_feat_map = train_outputs[layer_name][topk_indexes[t_idx]]\n",
        "                test_feat_map = test_outputs[layer_name][t_idx:t_idx + 1]\n",
        "                feat_gallery = topk_feat_map.transpose(3, 1).flatten(0, 2).unsqueeze(-1).unsqueeze(-1)\n",
        "                # calculate distance matrix\n",
        "                dist_matrix_list = []\n",
        "                for d_idx in range(feat_gallery.shape[0] // 100):\n",
        "                    dist_matrix = torch.pairwise_distance(feat_gallery[d_idx * 100:d_idx * 100 + 100], test_feat_map)\n",
        "                    dist_matrix_list.append(dist_matrix)\n",
        "                dist_matrix = torch.cat(dist_matrix_list, 0)\n",
        "                # k nearest features from the gallery (k=1)\n",
        "                score_map = torch.min(dist_matrix, dim=0)[0]\n",
        "                score_map = F.interpolate(score_map.unsqueeze(0).unsqueeze(0), size=224,\n",
        "                                          mode='bilinear', align_corners=False)\n",
        "                score_maps.append(score_map)\n",
        "\n",
        "            # average distance between the features\n",
        "            score_map = torch.mean(torch.cat(score_maps, 0), dim=0)\n",
        "\n",
        "            # apply gaussian smoothing on the score map\n",
        "            score_map = gaussian_filter(score_map.squeeze().cpu().detach().numpy(), sigma=4)\n",
        "            score_map_list.append(score_map)\n",
        "        flatten_gt_mask_list = np.concatenate(gt_mask_list).ravel()\n",
        "        flatten_score_map_list = np.concatenate(score_map_list).ravel()\n",
        "\n",
        "        # calculate per-pixel level ROCAUC\n",
        "        fpr, tpr, _ = roc_curve(flatten_gt_mask_list, flatten_score_map_list)\n",
        "        per_pixel_rocauc = roc_auc_score(flatten_gt_mask_list, flatten_score_map_list)\n",
        "        total_pixel_roc_auc.append(per_pixel_rocauc)\n",
        "        print('%s pixel ROCAUC: %.3f' % (class_name, per_pixel_rocauc))\n",
        "        fig_pixel_rocauc.plot(fpr, tpr, label='%s ROCAUC: %.3f' % (class_name, per_pixel_rocauc))\n",
        "\n",
        "        # get optimal threshold\n",
        "        precision, recall, thresholds = precision_recall_curve(flatten_gt_mask_list, flatten_score_map_list)\n",
        "        a = 2 * precision * recall\n",
        "        b = precision + recall\n",
        "        f1 = np.divide(a, b, out=np.zeros_like(a), where=b != 0)\n",
        "        threshold = thresholds[np.argmax(f1)]\n",
        "\n",
        "        # visualize localization result\n",
        "        visualize_loc_result(test_imgs, gt_mask_list, score_map_list, threshold, args.save_path, class_name, vis_num=5)\n",
        "    print('Average ROCAUC: %.3f' % np.mean(total_roc_auc))\n",
        "    fig_img_rocauc.title.set_text('Average image ROCAUC: %.3f' % np.mean(total_roc_auc))\n",
        "    fig_img_rocauc.legend(loc=\"lower right\")\n",
        "    print('Average pixel ROCUAC: %.3f' % np.mean(total_pixel_roc_auc))\n",
        "    fig_pixel_rocauc.title.set_text('Average pixel ROCAUC: %.3f' % np.mean(total_pixel_roc_auc))\n",
        "    fig_pixel_rocauc.legend(loc=\"lower right\")\n",
        "    fig.tight_layout()\n",
        "    fig.savefig(os.path.join(args.save_path, 'roc_curve.png'), dpi=100)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rflAKJhOTehP"
      },
      "source": [
        "def calc_dist_matrix(x, y):\n",
        "    \"\"\"Calculate Euclidean distance matrix with torch.tensor\"\"\"\n",
        "    n = x.size(0)\n",
        "    m = y.size(0)\n",
        "    d = x.size(1)\n",
        "    x = x.unsqueeze(1).expand(n, m, d)\n",
        "    y = y.unsqueeze(0).expand(n, m, d)\n",
        "    dist_matrix = torch.sqrt(torch.pow(x - y, 2).sum(2))\n",
        "    return dist_matrix"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OxEd3OCqTehQ"
      },
      "source": [
        "def visualize_loc_result(test_imgs, gt_mask_list, score_map_list, threshold,\n",
        "                         save_path, class_name, vis_num=5):\n",
        "    for t_idx in range(vis_num):\n",
        "        test_img = test_imgs[t_idx]\n",
        "        test_img = denormalization(test_img)\n",
        "        test_gt = gt_mask_list[t_idx].transpose(1, 2, 0).squeeze()\n",
        "        test_pred = score_map_list[t_idx]\n",
        "        test_pred[test_pred <= threshold] = 0\n",
        "        test_pred[test_pred > threshold] = 1\n",
        "        test_pred_img = test_img.copy()\n",
        "        test_pred_img[test_pred == 0] = 0\n",
        "        fig_img, ax_img = plt.subplots(1, 4, figsize=(12, 4))\n",
        "        fig_img.subplots_adjust(left=0, right=1, bottom=0, top=1)\n",
        "        for ax_i in ax_img:\n",
        "            ax_i.axes.xaxis.set_visible(False)\n",
        "            ax_i.axes.yaxis.set_visible(False)\n",
        "        ax_img[0].imshow(test_img)\n",
        "        ax_img[0].title.set_text('Image')\n",
        "        ax_img[1].imshow(test_gt, cmap='gray')\n",
        "        ax_img[1].title.set_text('GroundTruth')\n",
        "        ax_img[2].imshow(test_pred, cmap='gray')\n",
        "        ax_img[2].title.set_text('Predicted mask')\n",
        "        ax_img[3].imshow(test_pred_img)\n",
        "        ax_img[3].title.set_text('Predicted anomalous image')\n",
        "        os.makedirs(os.path.join(save_path, 'images'), exist_ok=True)\n",
        "        fig_img.savefig(os.path.join(save_path, 'images', '%s_%03d.png' % (class_name, t_idx)), dpi=100)\n",
        "        fig_img.clf()\n",
        "        plt.close(fig_img)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6qLn9PQyTehR"
      },
      "source": [
        "def denormalization(x):\n",
        "    mean = np.array([0.485, 0.456, 0.406])\n",
        "    std = np.array([0.229, 0.224, 0.225])\n",
        "    x = (((x.transpose(1, 2, 0) * std) + mean) * 255.).astype(np.uint8)\n",
        "    return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oNQ2g64ITehS"
      },
      "source": [
        "if __name__ == '__main__':\n",
        "    main()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}