create_tfrecord.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "import tensorflow as tf\n",
    "from PIL import Image\n",
    "from tensorflow.python.keras import backend as K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.enable_eager_execution()\n",
    "tf.executing_eagerly() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train_filenames = glob.glob(\"data/images/*.jpeg\")\n",
    "x_train_filenames.sort()\n",
    "\n",
    "y_train_filenames = glob.glob(\"data/masks/*.png\")\n",
    "y_train_filenames.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train_filenames, x_val_filenames, y_train_filenames, y_val_filenames = \\\n",
    "                    train_test_split(x_train_filenames, y_train_filenames, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = tf.data.Dataset.from_tensor_slices((x_train_filenames, y_train_filenames))\n",
    "val_ds = tf.data.Dataset.from_tensor_slices((x_val_filenames, y_val_filenames))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _bytes_feature(value):\n",
    "  \"\"\"Returns a bytes_list from a string / byte.\"\"\"\n",
    "  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def serialize_example(img,label_img):\n",
    "    \n",
    "    img_str = open(img,'rb').read()\n",
    "    label_img_str =open(label_img,'rb').read()\n",
    "    \n",
    "    feature = {\n",
    "        'image':_bytes_feature(img_str),\n",
    "        'mask':_bytes_feature(label_img_str)\n",
    "    }\n",
    "    \n",
    "    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))\n",
    "    return example_proto.SerializeToString()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tf_serialize_example(img, label_img):\n",
    "    tf_string = tf.py_func(\n",
    "        serialize_example, \n",
    "        (img, label_img), \n",
    "        tf.string)     \n",
    "    return tf.reshape(tf_string, ()) # The result is a scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = train_ds.map(tf_serialize_example)\n",
    "val_ds = val_ds.map(tf_serialize_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_tfrecord(filename, tfrecord):\n",
    "    filename = filename+'.tfrecord'\n",
    "    writer = tf.data.experimental.TFRecordWriter(filename)\n",
    "    writer.write(tfrecord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_tfrecord(\"train_ds\",train_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_tfrecord(\"val_ds\",train_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}