inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference notebook\n",
    "\n",
    "In this notebook, we make predictions in segmenting defects in hazelnut images from the MVTec Anomaly Detection [dataset](https://www.mvtec.com/company/research/datasets/mvtec-ad).\n",
    "\n",
    "\n",
    "First, import all the necessary libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "from processing.preprocessing import preprocess_data_from_images\n",
    "from processing.postprocessing import encode_masks_to_rgb\n",
    "from utils.dir_processing import clean_folder\n",
    "from utils.utils import create_color_map, create_category_dict\n",
    "\n",
    "from models import evaluate\n",
    "from models.model_functions import load_model\n",
    "from models.saving import load_params\n",
    "from models.train_test.test import test_model\n",
    "from models.saving import save_predictions\n",
    "\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All configuration parameters are defined in `env.yaml` file. These parameters include information about the location of the dataset and results folders, the architecture of the network, and its training parameters.\n",
    "\n",
    "Here we will use a pretrained model called `fuzzy_ivory_macaque`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load environment variables\n",
    "cfg = OmegaConf.load('configs/env.yaml')\n",
    "\n",
    "cfg.MODEL.model_name = \"whiny_red_mastiff\"\n",
    "\n",
    "# used for inference only\n",
    "categ_dict = create_category_dict(cfg.MODEL.categories)\n",
    "color_map = create_color_map(cfg.MODEL.categories)\n",
    "\n",
    "print('Class categories', categ_dict)\n",
    "\n",
    "num_classes = len(cfg.MODEL.categories) + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load model and weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use existing model with its proper name\n",
    "model_name = cfg.MODEL.model_name\n",
    "\n",
    "model = load_model(model_name)\n",
    "load_params(model_name, cfg)\n",
    "\n",
    "print(f'Model name is {model_name}')    \n",
    "    \n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we load and preprocess images and masks from the dataset. We split the initial dataset into training and testing subsets. \n",
    "\n",
    "We will use only the test subset for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === DATASET LOADING AND PREPROCESSING === #\n",
    "X, y = preprocess_data_from_images(data_path = cfg.DIRS.data, \n",
    "                                   shape = eval(cfg.DATA.img_dims),\n",
    "                                   categories = cfg.MODEL.categories)\n",
    "\n",
    "#=== TRAIN/TEST SPLIT === #\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y,\n",
    "                                                    test_size=cfg.DATA.test_split,\n",
    "                                                    shuffle=cfg.TRAINING.shuffle,\n",
    "                                                    random_state=cfg.DATA.seed)\n",
    "\n",
    "print(f'Number of TRAIN images: {len(X_train)}')\n",
    "print(f'Number of TEST images: {len(X_test)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# == INFERENCE == #\n",
    "y_pred = test_model(model, X_test, prediction_threshold=cfg.TRAINING.prediction_threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get confusion matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# == Confusion matrices == #\n",
    "confusion_classes, imgs_labels = evaluate.get_confusion_indices(y_test,\n",
    "                                                                y_pred,\n",
    "                                                                categories_dict = categ_dict,\n",
    "                                                                pixel_thres = cfg.TRAINING.pixel_threshold,\n",
    "                                                                meanIoU_threshold = cfg.TRAINING.iou_threshold)            \n",
    "            \n",
    "for class_name,confusion_matrix in confusion_classes.items():        \n",
    "    evaluate.save_confusion_matrix(confusion_matrix, model_name, class_name, class_counter=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save predictions sorting them into true positive (tp), true negative (tn), false positive (fp) and false negatives (fn) folders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encode ground truth and prediction masks\n",
    "y_test_en, y_pred_en = encode_masks_to_rgb(y_test, y_pred, color_map)\n",
    "\n",
    "PATH_RESULTS = os.path.join(cfg.DIRS.results, model_name)\n",
    "clean_folder(PATH_RESULTS)\n",
    "\n",
    "save_predictions(X_test, \n",
    "                 y_test_en,\n",
    "                 y_pred_en,\n",
    "                 PATH_RESULTS,\n",
    "                 imgs_labels,\n",
    "                 confusion_classes,\n",
    "                 color_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find confusion matrix as well as prediction masks in `results` folder."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "655936058c9faa9280730ea635086effdc23d6308265a7763ad2c8854fe885a3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}