MvTec_kaggle_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "language": "python",
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "version": "3.6.4",
      "file_extension": ".py",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "name": "python",
      "mimetype": "text/x-python"
    },
    "colab": {
      "name": "MvTec_kaggle_colab.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/AutoEncoder-SSIM-for-unsupervised-anomaly-detection-/blob/master/MvTec_kaggle_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lEV_7XY2unW9"
      },
      "source": [
        "# UNSUPERVISED ANOMALY SEGMENTATION\n",
        "### This is a implementation of **DFR: Deep Feature Reconstruction for Unsupervised Anomaly Segmentation** https://arxiv.org/abs/2012.07122\n",
        "#### We train only with OK images, and then we predict a hotmap, where high areas are presumably anomalous\n",
        "#### The pretrained weights provided by me (input anomalysegmentation-weights) are not very good because Kaggle limits training time to 9hrs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNbje4BzunW_"
      },
      "source": [
        "## IMPORTS"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T11:51:05.086186Z",
          "iopub.execute_input": "2021-10-27T11:51:05.08656Z",
          "iopub.status.idle": "2021-10-27T11:51:13.787564Z",
          "shell.execute_reply.started": "2021-10-27T11:51:05.086474Z",
          "shell.execute_reply": "2021-10-27T11:51:13.786715Z"
        },
        "trusted": true,
        "id": "Orrs_qg8unXA"
      },
      "source": [
        "!pip install tqdm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:21.069911Z",
          "iopub.execute_input": "2021-10-27T12:32:21.070465Z",
          "iopub.status.idle": "2021-10-27T12:32:27.585617Z",
          "shell.execute_reply.started": "2021-10-27T12:32:21.070422Z",
          "shell.execute_reply": "2021-10-27T12:32:27.58487Z"
        },
        "trusted": true,
        "id": "9-QxrFCyunXB"
      },
      "source": [
        "import tensorflow as tf  \n",
        "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.applications.vgg16 import VGG16\n",
        "from tensorflow.keras.applications.vgg19 import VGG19\n",
        "from tensorflow.keras import Input, Model, Sequential\n",
        "from tensorflow.keras import layers\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.decomposition import PCA\n",
        "\n",
        "from skimage import io\n",
        "from skimage.transform import resize\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "import matplotlib.pyplot as plt\n",
        "import cv2\n",
        "\n",
        "import random\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "\n",
        "from random import randrange"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fm8n_bHjunXC"
      },
      "source": [
        "## Custom Data Generator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:27.587408Z",
          "iopub.execute_input": "2021-10-27T12:32:27.587673Z",
          "iopub.status.idle": "2021-10-27T12:32:27.600578Z",
          "shell.execute_reply.started": "2021-10-27T12:32:27.587634Z",
          "shell.execute_reply": "2021-10-27T12:32:27.599697Z"
        },
        "trusted": true,
        "id": "vlE-wwJyunXC"
      },
      "source": [
        "class CustomDataGen(tf.keras.utils.Sequence):\n",
        "    \n",
        "    def __init__(self, path = \"../input/mvtecad-mvtec-anomaly-detection/mvtec_anomaly_detection/screw/\",\n",
        "                 batch_size = 4,\n",
        "                 input_size=(224, 224),\n",
        "                 shuffle=True, seed = None, subset = 'training'):\n",
        "        \n",
        "        self.image_data_generator = ImageDataGenerator(rescale=1. / 255, data_format='channels_last',     \n",
        "            #zoom_range = 0.1,\n",
        "            #width_shift_range = 0.05,\n",
        "            #height_shift_range = 0.05,\n",
        "            #brightness_range=(0.95,1.05)\n",
        "        )\n",
        "        \n",
        "        if seed is None:\n",
        "            random.randint(0, 2**32)\n",
        "            \n",
        "        self.batch_size = batch_size\n",
        "        self.input_size = input_size\n",
        "        \n",
        "        self.X_paths = [os.fsdecode(file) for file in os.scandir(f\"{path}/train/good\")]\n",
        "        self.X_train, self.X_val = train_test_split(self.X_paths, test_size=0.2, random_state = seed, shuffle = True)\n",
        "        \n",
        "        if subset == 'training':            \n",
        "            self.X = self.X_train\n",
        "        elif subset == 'validation':\n",
        "            self.X = self.X_val\n",
        "            \n",
        "        self.n = len(self.X)\n",
        "    \n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        \n",
        "        data_x = []\n",
        "        \n",
        "        for i in range(self.batch_size):\n",
        "            image = cv2.imread(self.X[self.batch_size * index + i])\n",
        "            \n",
        "            image = self.image_data_generator.random_transform(image)\n",
        "            \n",
        "            image = image / 255\n",
        "            \n",
        "            image = tf.image.resize(image, self.input_size)\n",
        "            data_x.append(image)\n",
        "            \n",
        "        data_x = np.array(data_x)\n",
        "        \n",
        "        return data_x, data_x\n",
        "    \n",
        "    def __len__(self):            \n",
        "        return self.n // self.batch_size"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:27.602198Z",
          "iopub.execute_input": "2021-10-27T12:32:27.60265Z",
          "iopub.status.idle": "2021-10-27T12:32:27.799395Z",
          "shell.execute_reply.started": "2021-10-27T12:32:27.602614Z",
          "shell.execute_reply": "2021-10-27T12:32:27.798596Z"
        },
        "trusted": true,
        "id": "0b3TJ20KunXD"
      },
      "source": [
        "INPUT_SIZE = (224,224)\n",
        "seed = random.randint(0, 2**32)\n",
        "train_datagen = CustomDataGen(input_size = INPUT_SIZE, batch_size = 4, seed = seed)\n",
        "validation_datagen = CustomDataGen(input_size = INPUT_SIZE, subset = \"validation\", seed = seed) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:27.802236Z",
          "iopub.execute_input": "2021-10-27T12:32:27.802966Z",
          "iopub.status.idle": "2021-10-27T12:32:27.811672Z",
          "shell.execute_reply.started": "2021-10-27T12:32:27.802927Z",
          "shell.execute_reply": "2021-10-27T12:32:27.810836Z"
        },
        "trusted": true,
        "id": "I1AVLhHfunXD"
      },
      "source": [
        "#Check we dont have the same images in train and validation sets\n",
        "any([value in train_datagen.X for value in validation_datagen.X])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:27.813225Z",
          "iopub.execute_input": "2021-10-27T12:32:27.813662Z",
          "iopub.status.idle": "2021-10-27T12:32:30.60739Z",
          "shell.execute_reply.started": "2021-10-27T12:32:27.813623Z",
          "shell.execute_reply": "2021-10-27T12:32:30.606679Z"
        },
        "trusted": true,
        "id": "j87g1P7YunXE"
      },
      "source": [
        "x, y = train_datagen[0]\n",
        "io.imshow(x[0])\n",
        "io.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RH2ig1AbunXE"
      },
      "source": [
        "## Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:36.822906Z",
          "iopub.execute_input": "2021-10-27T12:32:36.823328Z",
          "iopub.status.idle": "2021-10-27T12:32:36.889709Z",
          "shell.execute_reply.started": "2021-10-27T12:32:36.823291Z",
          "shell.execute_reply": "2021-10-27T12:32:36.88894Z"
        },
        "trusted": true,
        "id": "rvTQPx17unXE"
      },
      "source": [
        "input_model_filepath = \"../input/anomalysegmentation-weights/anomaly-segmentation-model_vgg19.h5\"\n",
        "output_model_filepath = \"./anomaly-segmentation-model.h5\"\n",
        "#model class\n",
        "#Implementation of DFR: Deep Feature Reconstruction for Unsupervised Anomaly Segmentation https://arxiv.org/abs/2012.07122\n",
        "class AnomalySegmentator(tf.keras.Model):\n",
        "    def __init__(self, init_layer = 0, end_layer = None):\n",
        "        super(AnomalySegmentator, self).__init__()\n",
        "        #self.L2_weight = 1e-6\n",
        "        self.init_layer = init_layer\n",
        "        self.end_layer = end_layer\n",
        "        \n",
        "    def build_autoencoder(self, c0, cd):\n",
        "        self.autoencoder = Sequential([\n",
        "            layers.InputLayer((self.map_shape[0]//4, self.map_shape[1]//4, c0)),\n",
        "            layers.Conv2D((c0 + cd) // 2,(1,1), padding='same', activation = tf.keras.layers.LeakyReLU(alpha=0.1)),\n",
        "            layers.Conv2D(2*cd,(1,1), padding='same', activation = tf.keras.layers.LeakyReLU(alpha=0.1)),\n",
        "            layers.Conv2D(cd,(1,1), padding='same'),\n",
        "            layers.Conv2D(2*cd,(1,1), padding='same', activation = tf.keras.layers.LeakyReLU(alpha=0.1)),\n",
        "            layers.Conv2D((c0 + cd) // 2,(1,1), padding='same', activation = tf.keras.layers.LeakyReLU(alpha=0.1)),\n",
        "            layers.Conv2D(c0,(1,1), padding='same')            \n",
        "            \n",
        "        ])\n",
        "        \n",
        "    def build(self, input_shape):\n",
        "        \n",
        "        self.vgg = VGG19(include_top = False, weights = 'imagenet', input_shape=input_shape[1:])\n",
        "        self.features_list = [layer.output for layer in self.vgg.layers if 'conv' in layer.name][self.init_layer:self.end_layer]\n",
        "        \n",
        "        self.feature_extractor = Model(inputs = self.vgg.input, \n",
        "                                       outputs = self.features_list)\n",
        "        self.feature_extractor.trainable = False   \n",
        "        \n",
        "        self.threshold = tf.Variable(0, trainable = False, dtype = tf.float32)\n",
        "\n",
        "        self.map_shape = self.features_list[0].shape[1:-1]\n",
        "        \n",
        "        self.average_pooling = layers.AveragePooling2D(pool_size=(4, 4), strides=(4,4))       \n",
        "        \n",
        "        \n",
        "        \n",
        "        self.c0 = sum([feature.shape[-1] for feature in self.features_list])        \n",
        "        self.cd = 40        \n",
        "        self.build_autoencoder(self.c0, self.cd)\n",
        "        \n",
        "          \n",
        "    def call(self, inputs):\n",
        "        features = self.feature_extractor(inputs)\n",
        "        resized_features = [tf.image.resize(feature, self.map_shape) for feature in features]\n",
        "        resized_features = tf.concat(resized_features, axis = -1)\n",
        "        \n",
        "        resized_features = self.average_pooling(resized_features)\n",
        "\n",
        "        autoencoder_output = self.autoencoder(resized_features)\n",
        "        return tf.reduce_sum((autoencoder_output - resized_features)**2, axis = -1)\n",
        "        \n",
        "    def reconstruction_loss(self):\n",
        "        @tf.function\n",
        "        def _loss(y_true, y_pred):\n",
        "            loss = tf.reduce_mean(y_pred, axis = (1,2)) / (tf.cast(tf.shape(y_pred)[0], tf.float32) * self.c0)\n",
        "            return loss\n",
        "                    \n",
        "        return _loss\n",
        "\n",
        "    def compute_threshold(self, data_loader, fpr = 0.05):\n",
        "      error = []\n",
        "      for i in tqdm(range(len(data_loader))):\n",
        "        x, y = data_loader[i]\n",
        "        error.append(self(x))\n",
        "\n",
        "      threshold = np.percentile(error, 100 - fpr)\n",
        "      self.threshold = tf.Variable(threshold, trainable = False, dtype = tf.float32)\n",
        "        \n",
        "    def compute_pca(self, data_loader):\n",
        "        extraction_per_sample = 20\n",
        "        \n",
        "        extractions = []        \n",
        "        for i in tqdm(range(len(data_loader))):\n",
        "            x, _ = data_loader[i]     \n",
        "            \n",
        "            features = self.feature_extractor(x)\n",
        "            resized_features = [tf.image.resize(feature, self.map_shape) for feature in features]\n",
        "            resized_features = tf.concat(resized_features, axis = -1)\n",
        "\n",
        "            resized_features = self.average_pooling(resized_features)\n",
        "            \n",
        "            for feature in resized_features:\n",
        "                \n",
        "                for _ in range(extraction_per_sample):                    \n",
        "                \n",
        "                    row, col = randrange(feature.shape[0]), randrange(feature.shape[1])\n",
        "                    extraction = feature[row, col]\n",
        "                    extractions.append(extraction)\n",
        "            \n",
        "        extractions = np.array(extractions)\n",
        "        print(f\"Extractions Shape: {extractions.shape}\")\n",
        "        pca = PCA(0.9, svd_solver = \"full\")\n",
        "        pca.fit(extractions)\n",
        "        self.cd = pca.n_components_\n",
        "        self.build_autoencoder(self.c0, self.cd)\n",
        "        print(f\"Components with explainable variance 0.9 -> {self.cd}\")\n",
        "        \n",
        "as_model = AnomalySegmentator()\n",
        "as_model.compile(Adam(1e-4), loss = as_model.reconstruction_loss())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:32:42.56686Z",
          "iopub.execute_input": "2021-10-27T12:32:42.56713Z",
          "iopub.status.idle": "2021-10-27T12:34:02.709736Z",
          "shell.execute_reply.started": "2021-10-27T12:32:42.567098Z",
          "shell.execute_reply": "2021-10-27T12:34:02.708967Z"
        },
        "trusted": true,
        "id": "8Qlp7j8UunXG"
      },
      "source": [
        "as_model.build((None, *INPUT_SIZE,3))\n",
        "as_model.compute_pca(train_datagen)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GIIw4UCYunXG"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T12:41:19.401357Z",
          "iopub.execute_input": "2021-10-27T12:41:19.402284Z"
        },
        "trusted": true,
        "id": "9yiZ11AuunXG"
      },
      "source": [
        "#True if we want to train the model\n",
        "if True:\n",
        "    # Training the model\n",
        "    plateau = tf.keras.callbacks.ReduceLROnPlateau(\n",
        "      monitor='val_loss', factor=0.5, patience=5, verbose = 1\n",
        "    )\n",
        "    es = tf.keras.callbacks.EarlyStopping(monitor='val_loss', mode='min', verbose=1,\n",
        "                                          patience=15)  # Early stopping (stops training when validation doesn't improve for {patience} epochs)\n",
        "    save_best = tf.keras.callbacks.ModelCheckpoint(output_model_filepath, monitor='val_loss', save_best_only=True,\n",
        "                                                mode='min', save_weights_only = True)  # Saves the best version of the model to disk (as measured on the validation data set)\n",
        "          \n",
        "    history = as_model.fit(train_datagen,\n",
        "        epochs=500,\n",
        "        validation_data=validation_datagen,\n",
        "        shuffle=True,\n",
        "        callbacks=[es, save_best, plateau])\n",
        "    #Training history\n",
        "    plt.plot(history.history['loss'])\n",
        "    plt.plot(history.history['val_loss'])\n",
        "    plt.title('model loss')\n",
        "    plt.ylabel('loss')\n",
        "    plt.xlabel('epoch')\n",
        "    plt.legend(['train', 'validation'], loc='upper left')\n",
        "    plt.show()\n",
        "    \n",
        "    as_model.load_weights(output_model_filepath)\n",
        "    as_model.compute_threshold(train_datagen)\n",
        "    as_model.save_weights(output_model_filepath)\n",
        "else:\n",
        "    \n",
        "    as_model.build((None, *INPUT_SIZE,3))\n",
        "    as_model.load_weights(input_model_filepath)\n",
        "    as_model.compute_threshold(train_datagen)\n",
        "    as_model.summary()\n",
        "    as_model.autoencoder.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y9nH0w7CunXH"
      },
      "source": [
        "# TESTING"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T11:52:21.816638Z",
          "iopub.execute_input": "2021-10-27T11:52:21.817352Z",
          "iopub.status.idle": "2021-10-27T11:52:21.93635Z",
          "shell.execute_reply.started": "2021-10-27T11:52:21.817294Z",
          "shell.execute_reply": "2021-10-27T11:52:21.935624Z"
        },
        "trusted": true,
        "id": "IfYMYJvJunXH"
      },
      "source": [
        "classes = ['good', 'manipulated_front', 'scratch_head', 'scratch_neck', 'thread_side', 'thread_top'] #only 0 index is GOOD\n",
        "test_datagen = ImageDataGenerator(rescale=1. / 255, data_format='channels_last')\n",
        "test_generator = test_datagen.flow_from_directory(\"../input/mvtecad-mvtec-anomaly-detection/mvtec_anomaly_detection/screw/test\", \n",
        "                                                  target_size = INPUT_SIZE,\n",
        "                                                  batch_size = 1,\n",
        "                                                  class_mode = 'sparse')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "execution": {
          "iopub.status.busy": "2021-10-27T11:52:21.939135Z",
          "iopub.execute_input": "2021-10-27T11:52:21.939522Z",
          "iopub.status.idle": "2021-10-27T11:52:29.881655Z",
          "shell.execute_reply.started": "2021-10-27T11:52:21.939486Z",
          "shell.execute_reply": "2021-10-27T11:52:29.880975Z"
        },
        "trusted": true,
        "id": "yfV9qkbvunXH"
      },
      "source": [
        "for i in range(10):\n",
        "    x_batch, y_batch = next(test_generator)    \n",
        "    hotmaps = as_model(x_batch)\n",
        "    for x, y, hotmap in zip(x_batch, y_batch, hotmaps):\n",
        "\n",
        "        prediction = np.any(hotmap > as_model.threshold)\n",
        "\n",
        "        hotmap = resize(hotmap, x.shape[:-1], anti_aliasing = True)\n",
        "        mask = np.where(hotmap > as_model.threshold, 1, 0)\n",
        "        \n",
        "\n",
        "        f, axarr = plt.subplots(1,2, figsize=(15,15))\n",
        "        axarr[0].imshow(x)\n",
        "        axarr[1].imshow(mask)\n",
        "        axarr[1].imshow(x, alpha = 0.75)\n",
        "        plt.show()\n",
        "\n",
        "        print(f\"Threshold: {as_model.threshold.numpy()} MaxValue: {hotmap.max()}\")\n",
        "        print(f\"GT: {classes[int(y)]}, anomaly detected: {prediction}\")\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}