验证码的tfrecord文件生成.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import os\n",
    "import random\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">> converting image 6216/6216\n",
      ">> converting image 100/100\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "_RANDOM_SEED=0\n",
    "NUM_TEST=100\n",
    "DATASET_DIR='D:/Tensorflow/captcha/images/'\n",
    "TFRECORD_DIR='D:/Tensorflow/captcha/'\n",
    "def data_exists(dataset_dir):\n",
    "    for split_name in ['train','test']:\n",
    "        output_filename=os.path.join(dataset_dir,split_name+'.tfrecords')\n",
    "        if not tf.gfile.Exists(output_filename):\n",
    "            return False\n",
    "    return True\n",
    "def get_file_and_classnames(dataset_dir):\n",
    "    photo_filenames=[]\n",
    "    for filename in os.listdir(dataset_dir):\n",
    "        path=os.path.join(dataset_dir,filename)\n",
    "        photo_filenames.append(path)    \n",
    "    return photo_filenames\n",
    "def int64_feature(values):\n",
    "    if not isinstance(values,(tuple,list)):\n",
    "        values=[values]\n",
    "    return tf.train.Feature(int64_list=tf.train.Int64List(value=values))\n",
    "\n",
    "def bytes_feature(values):\n",
    "    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))\n",
    "def image_to_tfexample(image_data,labels):\n",
    "    return tf.train.Example(features=tf.train.Features(feature={\n",
    "        'image':bytes_feature(image_data),\n",
    "        'label0': int64_feature(labels[0]),\n",
    "        'label1': int64_feature(labels[1]),\n",
    "        'label2': int64_feature(labels[2]),\n",
    "        'label3': int64_feature(labels[3]),\n",
    "    }))\n",
    "def _convert_dataset(split_name,filenames,dataset_dir):\n",
    "    assert split_name in ['train','test']\n",
    "    with tf.Session() as sess:\n",
    "        output_filename=os.path.join(TFRECORD_DIR,split_name+'.tfrecords')\n",
    "        with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:\n",
    "            for i,filename in enumerate(filenames):\n",
    "                try:\n",
    "                    sys.stdout.write('\\r>> converting image %d/%d'% (i+1,len(filenames)))\n",
    "                    sys.stdout.flush()\n",
    "                    image_data=Image.open(filename)\n",
    "                    image_data=image_data.resize((224,224))\n",
    "                    image_data=np.array(image_data.convert('L'))\n",
    "                    image_data=image_data.tobytes()\n",
    "                    labels=filename.split('/')[-1][0:4]\n",
    "                    num_labels=[]\n",
    "                    for j in range(4):\n",
    "                        num_labels.append(int(labels[j]))\n",
    "                    example=image_to_tfexample(image_data,num_labels)\n",
    "                    tfrecord_writer.write(example.SerializeToString())\n",
    "                except IOError as e:\n",
    "                    print('could not read:',filename)\n",
    "                    print('Error :',e)\n",
    "                    print('skip it\\n')\n",
    "        sys.stdout.write('\\n')\n",
    "        sys.stdout.flush()           \n",
    "                    \n",
    "                    \n",
    "if __name__=='__main__':\n",
    "    if data_exists(TFRECORD_DIR):\n",
    "        print('file exsits')\n",
    "    else:\n",
    "        photo_filenames=get_file_and_classnames(DATASET_DIR)\n",
    "        random.seed(_RANDOM_SEED)\n",
    "        random.shuffle(photo_filenames)\n",
    "        training_filenames=photo_filenames[NUM_TEST:]\n",
    "        testing_filenames=photo_filenames[:NUM_TEST]\n",
    "        _convert_dataset('train',training_filenames,DATASET_DIR)\n",
    "        _convert_dataset('test',testing_filenames,DATASET_DIR)\n",
    "        print('done')\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}