deep_learning/space_image/Classification_Inference_from_PB_Tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update && apt install -y libsm6 libxext6 libfontconfig1 libxrender1\n",
    "!pip install opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import tensorflow as tf\n",
    "from collections import OrderedDict\n",
    "import re\n",
    "from PIL import Image\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import cv2\n",
    "\n",
    "print(\"Tensorflow version: {}\".format(tf.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model and Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "PB: the frozen model\n",
    "pb.txt: the label map file\n",
    "\"\"\"\n",
    "\n",
    "PATH_TO_PB = \"/notebooks/model/inception_single_0.1/inception_single_0.1.pb\"\n",
    "PATH_TO_LABELS = \"/notebooks/model/inception_single_0.1/image_label.pbtxt\"\n",
    "\n",
    "if not os.path.exists(PATH_TO_PB): raise FileNotFoundError(\"PB is not found.\")\n",
    "if not os.path.exists(PATH_TO_LABELS): raise FileNotFoundError(\"Label is not found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## load frozen graph "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detection_graph = tf.Graph()\n",
    "with detection_graph.as_default():\n",
    "    od_graph_def = tf.GraphDef()\n",
    "    with tf.gfile.GFile(PATH_TO_PB, 'rb') as fid:\n",
    "        serialized_graph = fid.read()\n",
    "        od_graph_def.ParseFromString(serialized_graph)\n",
    "        tf.import_graph_def(od_graph_def, name='')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show all operation names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_operation_names(count=10):\n",
    "    with detection_graph.as_default():\n",
    "        with tf.Session() as sess:\n",
    "            opts = tf.get_default_graph().get_operations()\n",
    "            for opt in opts[:count]: \n",
    "                for output in opt.outputs: print(output.name)\n",
    "            print(\"...\")\n",
    "            for opt in opts[-count:]: \n",
    "                for output in opt.outputs: print(output.name)\n",
    "                    \n",
    "show_operation_names(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Label file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "category_index = OrderedDict()\n",
    "name_indx = OrderedDict()\n",
    "count = 1\n",
    "with open(PATH_TO_LABELS, \"r\") as fin:\n",
    "    tmpData = \"\"\n",
    "    for line in fin:\n",
    "        tmpData = line.strip()\n",
    "        category_index[count] = tmpData\n",
    "        name_indx[tmpData] = count\n",
    "        count += 1\n",
    "    print(category_index)\n",
    "    print(name_indx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference_single_image(image, graph):\n",
    "    with graph.as_default():\n",
    "        with tf.Session() as sess:\n",
    "            # handle input and output tensor\n",
    "            opts = tf.get_default_graph().get_operations()\n",
    "            all_tensorflow_names = { output.name for opt in opts for output in opt.outputs }\n",
    "            tensor_dict = {}\n",
    "            for key in ['final_result']:\n",
    "                tensor_name = key + ':0'\n",
    "                if tensor_name in all_tensorflow_names:\n",
    "                    tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(tensor_name)\n",
    "                    \n",
    "            # run for single image            \n",
    "            # input\n",
    "            image_tensor = tf.get_default_graph().get_tensor_by_name('Placeholder:0')\n",
    "            \n",
    "            # inference\n",
    "            output_dict = sess.run(tensor_dict, feed_dict={image_tensor: np.expand_dims(image, 0)})\n",
    "            \n",
    "            # convert data type float32 to appropriate\n",
    "            output_dict['final_result'] = output_dict['final_result']\n",
    "            \n",
    "        return output_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_image(imagePath):\n",
    "    image_path = imagePath\n",
    "    if not os.path.exists(image_path): raise FileNotFoundError(\"{} not found.\".format(image_path))\n",
    "        \n",
    "    image = cv2.imread(image_path)\n",
    "    image = cv2.resize(image, (299, 299), interpolation=cv2.INTER_CUBIC)\n",
    "    image = image / 255.\n",
    "    image_np = image[:,:,::-1]\n",
    "    output_dict = inference_single_image(image_np, detection_graph)\n",
    "    \n",
    "    return output_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = '/notebooks/imgs/image1.jpg'\n",
    "output_dict = single_image(image_path)\n",
    "cls_idx = int(np.argmax(output_dict['final_result'], axis=1) + 1)\n",
    "print(\"Classification: {}\".format(category_index[cls_idx]))"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}