BaseImageData.py 文件源码

def write_tfrecord(self, img_list, label_list, record_path):
        # write a single tfrecord
        if os.path.exists(record_path):
            print ("%s exists!"%record_path)
            return

        self._check_list()
        print ("write %s"%record_path)
        self._write_info()

        writer = tf.python_io.TFRecordWriter(record_path)
        c = 0
        for imgname,label in zip(img_list,label_list):

            img = Image.open(imgname).resize((self.flags.width, self.flags.height))
            data = np.array(img).astype(np.uint8)
            img,data = self._check_color(img,data)

            example = self._get_example(data,label)
            writer.write(example.SerializeToString())
            c+=1
            if c%LOG_EVERY == 0:
                print ("%d images written to tfrecord"%c)
        writer.close()
        print("writing %s done"%record_path)