人脸识别原理详解：使用tfrecord集合网络训练所有数据

def  get_dataset(dir, item): #对应生成的train_pnet_landmark.txt
    print('dir is: ', dir, item)
    dataset_dir = os.path.join(dir, item)
    print('join path: ', os.path.join(dir, item))
    print('dataset dir: ', dataset_dir)
    image_list = open(dataset_dir, 'r')
    dataset = []
    for line in tqdm(image_list.readlines()):
        info = line.strip().split(' ')
        if len(info) < 2:
            print('info err: ', info)

        data_example = {}
        bbox = {}
        data_example['filename'] = info[0]
        data_example['label'] = int(info[1])
        bbox['xmin'] = 0.0  #初始化人脸区域
        bbox['ymin'] = 0.0
        bbox['xmax'] = 0.0
        bbox['ymax'] = 0.0
        bbox['xlefteye'] = 0.0 #初始化10个关键点
        bbox['ylefteye'] = 0.0
        bbox['xrighteye'] = 0.0
        bbox['yrighteye'] = 0.0
        bbox['xnose'] = 0.0
        bbox['ynose'] = 0.0
        bbox['xleftmouth'] = 0.0
        bbox['yleftmouth'] = 0.0
        bbox['xrightmouth'] = 0.0
        bbox['yrightmouth'] = 0.0
        if len(info) == 6: #当前记录只包含人脸区域
            bbox['xmin'] = float(info[2])
            bbox['ymin'] = float(info[3])
            bbox['xmax'] = float(info[4])
            bbox['ymax'] = float(info[5])
        if len(info) == 12:  #当前记录包含了10个人脸关键点
            bbox['xlefteye'] =  float(info[2])#初始化10个关键点
            bbox['ylefteye'] = float(info[3])
            bbox['xrighteye'] = float(info[4])
            bbox['yrighteye'] = float(info[5])
            bbox['xnose'] = float(info[6])
            bbox['ynose'] = float(info[7])
            bbox['xleftmouth'] = float(info[8])
            bbox['yleftmouth'] = float(info[9])
            bbox['xrightmouth'] = float(info[10])
            bbox['yrightmouth'] = float(info[11])
        data_example['bbox'] = bbox
        dataset.append(data_example)
    return dataset

def  process_image(filename):
    try:
      image = cv2.imread(filename)
      #print('image to string')
      image_data = image.tostring()
      #print('finish image to string')
      assert len(image.shape) == 3
      height = image.shape[0]
      width = image.shape[1]
      assert image.shape[2] == 3
      return image_data, height, width
    except Exception as e:
      #print('process image err: ', e, filename)
      return None, None, None

def  _int64_feature(value):
    if not isinstance(value, list):
        value = [value]
    try:
        return tf.train.Feature(int64_list=tf.train.Int64List(value = value))
    except Exception as e:
        print('int64 err: ', e)
def  _float_feature(value):
    if not isinstance(value, list):
        value = [value]
    try:
        return tf.train.Feature(float_list=tf.train.FloatList(value = value))
    except Exception as e:
        print('float err: ', e)
def  _bytes_feature(value):
    if not isinstance(value, list):
        value = [value]
    try:
        return tf.train.Feature(bytes_list = tf.train.BytesList(value = value))
    except Exception as e:
        print('bytes err: ', e)

def  convert_to_example(image_example, image_buffer):
    class_label = image_example['label']
    bbox = image_example['bbox']
    roi = [bbox['xmin'], bbox['ymin'], bbox['xmax'], bbox['ymax']]
    landmark = [bbox['xlefteye'], bbox['ylefteye'], bbox['xrighteye'], bbox['yrighteye'],
                bbox['xnose'], bbox['ynose'], bbox['xleftmouth'],
                bbox['yleftmouth'], bbox['xrightmouth'], 
                bbox['yrightmouth']]
    try:
      example = tf.train.Example(features = tf.train.Features(feature = {
          'image/encoded': _bytes_feature(image_buffer),
          'image/label': _int64_feature(class_label),
          'image/roi': _float_feature(roi),
          'image/landmark': _float_feature(landmark)
      }))
      return example
    except Exception as e:
      print('example err: ',e , image_example)

def  add_to_tfrecord(filename, image_example, tfrecord_writer):
    begin = time.time()
    image_data, height, width = process_image(filename)
    end = time.time() 
   # print('time for process image: ', end-begin, filename)
    if image_data != None:
       # print('convert to example ')
        example = convert_to_example(image_example, image_data)
       # print('after convert to example')
        tfrecord_writer.write(example.SerializeToString())
        print('tfrecord write')

dataset_dir = '/content/drive/MyDrive/my_mtcnn/data'
import tensorflow as tf
def  create_tf_record(size):
    output_dir = os.path.join(dataset_dir, str(size) + "/tf_record")
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    if size == 12:
        net = 'PNet'
        tf_filenames = [os.path.join(output_dir, 'train%s_landmark.tfrecord' %(net))]
        items = ['train_pnet_landmark.txt']
    elif size == 24: #以后再考虑余下两只情况
        pass  
    elif  size == 48:
        pass
    if  tf.io.gfile.exists(tf_filenames[0]):
        print("tf record file alreay created")
    for tf_filename, item in zip(tf_filenames, items): #在size=12时看似多于，在后面处理size=24或48时用上
        print('reading daa....')
        dataset = get_dataset(dataset_dir, item)
        tf_filename = tf_filename + '_shuffle'
        random.shuffle(dataset)
        print('transform to tfrecord')
        with tf.io.TFRecordWriter(tf_filename) as tfrecord_writer:
            for image_example in tqdm(dataset):
                filename = image_example['filename']
                try:
                    add_to_tfrecord(filename, image_example, tfrecord_writer)
                except Exception as e:
                    print('tf record exception: ', e)
    print('completing transform..!')

create_tf_record(12)

  example = tf.train.Example(features = tf.train.Features(feature = {
          'image/encoded': _bytes_feature(image_buffer),
          'image/label': _int64_feature(class_label),
          'image/roi': _float_feature(roi),
          'image/landmark': _float_feature(landmark)
      }))