KNN-Classifier-Steel-Surface-Defects.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KNN Classifier for Steel surface defect recognition\n",
    "\n",
    "#### If you use this code, please cite our work: [[paper](https://doi.org/10.1016/j.promfg.2020.05.146)][[code](https://github.com/adipandas/one-shot-steel-surfaces)]\n",
    "\n",
    "### Aditya M. Deshpande, Ali A. Minai, Manish Kumar, One-Shot Recognition of Manufacturing Defects in Steel Surfaces, Procedia Manufacturing, Volume 48, 2020, Pages 1064-1071, ISSN 2351-9789, https://doi.org/10.1016/j.promfg.2020.05.146.\n",
    "\n",
    "#### MIT License\n",
    "\n",
    "Copyright (c) 2020 Aditya M. Deshpande\n",
    "\n",
    "Permission is hereby granted, free of charge, to any person obtaining a copy\n",
    "of this software and associated documentation files (the \"Software\"), to deal\n",
    "in the Software without restriction, including without limitation the rights\n",
    "to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
    "copies of the Software, and to permit persons to whom the Software is\n",
    "furnished to do so, subject to the following conditions:\n",
    "\n",
    "The above copyright notice and this permission notice shall be included in all\n",
    "copies or substantial portions of the Software.\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
    "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
    "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
    "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
    "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
    "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
    "SOFTWARE.\n",
    "\n",
    "### Note: for this case, the number of neighbors in KNN were set to `1` and the dataset was also configured accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.model_selection import train_test_split\n",
    "from imutils import paths\n",
    "import numpy as np\n",
    "import argparse\n",
    "import imutils\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### KNN classifier arguments\n",
    "\n",
    "#### NOTE: You may need to change the path in the below dictionary as per your dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = {'training_dataset': './dataset_all/knn_data/training',\n",
    "        'testing_dataset': './dataset_all/knn_data/testing',\n",
    "        'neighbors': 1,\n",
    "        'jobs': 1}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert image to Vector\n",
    "Function to *resize* the image to a fixed size, then flatten the image into a list of raw pixel intensities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_to_feature_vector(image, size=(32, 32)):\n",
    "    return cv2.resize(image, size).flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract a 3D color histogram from the HSV color space using the supplied number of `bins` per channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_color_histogram(image, bins=(8, 8, 8)):\n",
    "    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)\n",
    "    hist = cv2.calcHist([hsv], [0, 1, 2], None, bins, [0, 180, 0, 256, 0, 256])\n",
    "    cv2.normalize(hist, hist)\n",
    "    return hist.flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "imagePaths = list(paths.list_images(args[\"training_dataset\"]))\n",
    "imagePaths_test = list(paths.list_images(args[\"testing_dataset\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper function to prepare data\n",
    "* Extract each image and corresponding class label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_data(imagePaths):\n",
    "    rawImages, features, labels = [], [], []\n",
    "    \n",
    "    for (i, imagePath) in enumerate(imagePaths):    \n",
    "        image = cv2.imread(imagePath)\n",
    "        label = imagePath.split(os.path.sep)[-1].split(\".\")[0].split(\"_\")[0]\n",
    "        pixels = image_to_feature_vector(image)\n",
    "        hist = extract_color_histogram(image)\n",
    "        \n",
    "        rawImages.append(pixels)\n",
    "        features.append(hist)\n",
    "        labels.append(label)\n",
    "    \n",
    "    return rawImages, features, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load training and testing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images, train_features, train_labels = prepare_data(imagePaths)\n",
    "test_images, test_features, test_labels = prepare_data(imagePaths_test)\n",
    "\n",
    "print(\"[INFO] Shape of the image sampled from dataset:\", train_images[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Memory consumed by the raw images matrix and features matrix\n",
    "\n",
    "##### Training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainRI = np.array(train_images)\n",
    "trainRL = np.array(train_labels)\n",
    "trainF = np.array(train_features)\n",
    "print(\"[INFO] pixels matrix: {:.2f}MB\".format(trainRI.nbytes / (1024 * 1000.0)))\n",
    "print(\"[INFO] features matrix: {:.2f}MB\".format(trainF.nbytes / (1024 * 1000.0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Testing set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testRI = np.array(test_images)\n",
    "testRL = np.array(test_labels)\n",
    "testF = np.array(test_features)\n",
    "print(\"[INFO] pixels matrix: {:.2f}MB\".format(testRI.nbytes / (1024 * 1000.0)))\n",
    "print(\"[INFO] features matrix: {:.2f}MB\".format(testF.nbytes / (1024 * 1000.0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train and evaluate the k-NN classifer on the raw pixel intensities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print(\"[INFO] evaluating raw pixel accuracy...\")\n",
    "model = KNeighborsClassifier(n_neighbors=args[\"neighbors\"], n_jobs=args[\"jobs\"])\n",
    "model.fit(trainRI, trainRL)\n",
    "acc = model.score(testRI, testRL)\n",
    "\n",
    "print(\"[INFO] raw pixel accuracy: {:.2f}%\".format(acc * 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train and test KNN-classifier on Histogram Representation of image data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_test_data = train_test_split(features, labels, test_size=0.25, random_state=42)\n",
    "\n",
    "train_features, test_features, trainLabels, testLabels = train_test_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Histogram representations - train and evaluate k-NN classifer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"[INFO] evaluating histogram accuracy...\")\n",
    "\n",
    "model = KNeighborsClassifier(n_neighbors=args[\"neighbors\"],n_jobs=args[\"jobs\"])\n",
    "\n",
    "model.fit(train_features, trainLabels)\n",
    "\n",
    "acc = model.score(test_features, testLabels)\n",
    "print(\"[INFO] histogram accuracy: {:.2f}%\".format(acc * 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}