Visualize a TFRecord for Tensorflow Object Detection Library

visualise_tfrecord.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import sys\n",
    "\n",
    "import IPython.display\n",
    "import PIL\n",
    "\n",
    "MODELS_BASE = '/home/ubuntu/models/research'\n",
    "sys.path.append(MODELS_BASE)\n",
    "sys.path.append(MODELS_BASE + '/object_detection')\n",
    "sys.path.append(MODELS_BASE + '/slim')\n",
    "\n",
    "from object_detection.utils import visualization_utils as vu\n",
    "from object_detection.protos import string_int_label_map_pb2 as pb\n",
    "from object_detection.data_decoders.tf_example_decoder import TfExampleDecoder as TfDecoder\n",
    "from google.protobuf import text_format \n",
    "import itertools\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualise(tfrecords_filename, label_map=None):\n",
    "    if label_map is not None:\n",
    "        label_map_proto = pb.StringIntLabelMap()\n",
    "        with tf.gfile.GFile(label_map,'r') as f:\n",
    "            text_format.Merge(f.read(), label_map_proto)\n",
    "            class_dict = {}\n",
    "            for entry in label_map_proto.item:\n",
    "                class_dict[entry.id] = {'name':entry.display_name}\n",
    "    sess = tf.Session()\n",
    "    decoder = TfDecoder(label_map_proto_file=label_map, use_display_name=False)\n",
    "    sess.run(tf.tables_initializer())\n",
    "    topN = itertools.islice(tf.python_io.tf_record_iterator(tfrecords_filename), 5)\n",
    "    for record in topN:\n",
    "        example = decoder.decode(record)\n",
    "        host_example = sess.run(example)\n",
    "        scores = np.ones(host_example['groundtruth_boxes'].shape[0])\n",
    "        vu.visualize_boxes_and_labels_on_image_array( \n",
    "            host_example['image'],                                               \n",
    "            host_example['groundtruth_boxes'],                                                     \n",
    "            host_example['groundtruth_classes'],\n",
    "            scores,\n",
    "            class_dict,\n",
    "            max_boxes_to_draw=None,\n",
    "            use_normalized_coordinates=True)\n",
    "        \n",
    "        IPython.display.display(PIL.Image.fromarray(host_example['image']))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualise(\"train.record\",\"label_map.pbtxt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_tensorflow_p36)",
   "language": "python",
   "name": "conda_tensorflow_p36"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

Modified from the original code at https://stackoverflow.com/q/50391967
Thanks to Steve Goley for adding it there!