当前位置:网站首页>Convert verification code image to tfrecord file

Convert verification code image to tfrecord file

2022-06-26 08:40:00 Algorithmic Pilgrim

import tensorflow as tf
import os

"""  Be careful :  This API The processed image format is : The name of the picture is the content of the label ,  For example, a picture is called '2a3w.jpg', 2a3w Is the tag value .  The picture ,  The label is converted to TFrecord Format ,  Where the label has been [2a3w] >> [12, 23, 45, 46] """

FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string("pic_dir", "../pic_source/train6000/", " Path to source picture ")
tf.flags.DEFINE_string("letter", "abcdefghijklmnopqrstuvwxyz1234567890", " Type of character verification code ")
tf.flags.DEFINE_string("tfrecords_dir", "../pic_source/train6000_tfrecord/train6000.tfrecords", " Verification Code tfrecords file ")
tf.flags.DEFINE_integer("image_num",367," How many pictures are there in total ")

def dealWithLabel(label_str):
    """ :param label_str: ['2a2s', '2w3e', '4e5r', '3e5r', '4r5g'....] :return: """
    #  Build character index  {0:'a', 1:'b'......}
    num2letter = dict(enumerate(list(FLAGS.letter)))
    #  Key value pairs flip  {'a':0, 'b':1......}
    letter2num = dict(zip(num2letter.values(),num2letter.keys()))
    print(letter2num)

    #  Build a list of tags 
    array = []
    for lablex in label_str:

        letter_list = []

        for letter in list(str(lablex)): # '2a2s'>>[2,a,2,s]
            letter2_num = letter2num[str(letter)] #a>>2
            letter_list.append(letter2_num)

        array.append(letter_list)
    print(array)

    #  take array convert to tensor type 
    lable = tf.constant(array)

    return lable

def get_image(file_name):
    """  Get the picture data of the verification code , And labels   Be careful :  This API The processed image format is : The name of the picture is the content of the label , For example, a picture is called  '2a3w.jpg',2a3w Is the tag value . :return: image_batch,lable_batch """

    #  Traverse to get tag name  ['2a2s', '2w3e', '4e5r', '3e5r', '4r5g'....]
    # file_name = [str(lable).split(".")[0] for lable in os.listdir(path=FLAGS.pic_dir)]
    #  Construction path + file  ['../pic_source/train6000/2a2s.jpg', '', '', '', '', ''.....]
    file_list = [os.path.join(FLAGS.pic_dir, labl + ".jpg") for labl in file_name]
    #  Construct file queues 
    image_queue = tf.train.string_input_producer(file_list, shuffle=False)
    #  Build a reader 
    image_reader = tf.WholeFileReader()
    #  Read the picture content 
    key, value = image_reader.read(image_queue)
    #  Decode picture data 
    image = tf.image.decode_jpeg(value)
    image.set_shape([60, 180, 3]) #  Must follow height high  * weight wide  *  The channel number   Write... In the order of 
    #  Batch data  [367, 60, 180, 3]
    image_batch = tf.train.batch([image], batch_size=len(file_name),num_threads=1, capacity=len(file_name))
    # lable_batch = tf.train.batch(lable_queue, batch_size=len(file_name), num_threads=1, capacity=len(file_name))

    print(image_batch)
    # print(lable_batch)
    return image_batch

def write2tfrecords(image_batch, lable_batch):
    """  Write picture content and labels to tfrecords The file of  :param image_batch: The eigenvalue  :param lable_batch: Label value  :return: """
    #  Type conversion 
    lable_batch = tf.cast(lable_batch,tf.uint8)
    print(lable_batch)

    #  establish TFRecords Memory 
    writer = tf.python_io.TFRecordWriter(FLAGS.tfrecords_dir)

    #  Loop to construct the data on each picture example Protocol block , Write... After serialization 
    for i in range(FLAGS.image_num):
        #  Take out No i Picture data , Convert the corresponding type , The characteristic value of the picture should be converted into string form 
        image_string = image_batch[i].eval().tostring()

        #  Label value , Convert to integer 
        label_string = lable_batch[i].eval().tostring()

        #  Construct protocol block 
        example = tf.train.Example(features=tf.train.Features(feature={
    
            "image": tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_string])),
            "label": tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_string]))
        }))

        writer.write(example.SerializeToString())
    #  Close file 
    writer.close()

    return None

def generateTFrecord(file_name):
    """  Be careful :  This API The processed image format is : The name of the picture is the content of the label , For example, a picture is called  '2a3w.jpg',2a3w Is the tag value .  The picture , The label is converted to TFrecord Format , Where the label has been  [2a3w]>>[12,23,45,46] :param file_name: [] :return: """
    #  Get the picture in the verification code 
    image_batch = get_image(file_name)
    #  Get tag data 
    label_batch = dealWithLabel(label_str=file_name)
    print(image_batch, label_batch)

    with tf.Session() as sess:
        coord = tf.train.Coordinator() # Build queue 
        threads = tf.train.start_queue_runners(sess=sess, coord=coord) # Build thread 

        #  Write picture data and content to tfrecords The file of 
        write2tfrecords(image_batch, label_batch)
        coord.request_stop()
        coord.join(threads)
        print(' Generate TFRecord File successfully .............................................................')
    return None

if __name__ == '__main__':
    file_name = [str(lable).split(".")[0] for lable in os.listdir(path=FLAGS.pic_dir)]
    generateTFrecord(file_name=file_name)

原网站

版权声明
本文为[Algorithmic Pilgrim]所创,转载请带上原文链接,感谢
https://yzsam.com/2022/02/202202170555218733.html