xuelunshen / h-net Goto Github PK

import tensorflow as tf
import cv2
import sys
import numpy as np
from random import shuffle
import glob
import os
shuffle_data = True 


addrs = []

DATADIR_PHOTO = r"X:\xx\xx\real" 
DATADIR_UDS = r"X:\xx\xx\simulated"

for (photo, uds) in zip(os.listdir(DATADIR_PHOTO), os.listdir(DATADIR_UDS)):
        addrs.append([os.path.join(DATADIR_PHOTO, photo), os.path.join(DATADIR_UDS,uds)])
   
labels = [-1 if 'real' in addr else -1 for addr in addrs]

# to shuffle data
if shuffle_data:
    c = list(zip(addrs, labels))
    shuffle(c)
    addrs, labels = zip(*c)

# Divide the hata into 60% train, 20% validation, and 20% test
train_addrs = addrs[0:int(0.6*len(addrs))]
train_labels = labels[0:int(0.6*len(labels))]

val_addrs = addrs[int(0.6*len(addrs)):int(0.8*len(addrs))]
val_labels = labels[int(0.6*len(addrs)):int(0.8*len(addrs))]

test_addrs = addrs[int(0.8*len(addrs)):]
test_labels = labels[int(0.8*len(labels)):]    


def load_image(addr):
    # read an image and resize to (224, 224)
    # cv2 load images as BGR, convert it to RGB
    img = cv2.imread(addr)
    img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_CUBIC)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.uint8)
    return img


def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

train_filename = 'train.tfrecords'  # address to save the TFRecords file
# open the TFRecords file
writer = tf.python_io.TFRecordWriter(train_filename)

for i in range(len(train_addrs)):
    
    # print how many images are saved every 1000 images
    if not i % 1000:
        print ('Train data: {}/{}'.format(i, len(train_addrs)))
        sys.stdout.flush()
    # Load the image
    photo = load_image(train_addrs[i][0])
    uds = load_image(train_addrs[i][1])
    label = train_labels[i]
     
    # Create a feature
    feature = {'label': _int64_feature(label),
               'photo': _bytes_feature(tf.compat.as_bytes(photo.tostring())),
               'uds': _bytes_feature(tf.compat.as_bytes(uds.tostring()))}
   
    # Create an example protocol buffer
    example = tf.train.Example(features=tf.train.Features(feature=feature))
    
    # Serialize to string and write on the file
    writer.write(example.SerializeToString())
    
writer.close()
sys.stdout.flush()

# open the TFRecords file
val_filename = 'val.tfrecords'  # address to save the TFRecords file
writer = tf.python_io.TFRecordWriter(val_filename)
for i in range(len(val_addrs)):
    # print how many images are saved every 1000 images
    if not i % 1000:
        print ('Val data: {}/{}'.format(i, len(val_addrs)))
        sys.stdout.flush()
    # Load the image
    photo = load_image(val_addrs[i][0])
    uds = load_image(val_addrs[i][1])
    label = val_labels[i]
    # Create a feature
    feature = {'label': _int64_feature(label),
               'photo': _bytes_feature(tf.compat.as_bytes(photo.tostring())),
               'uds': _bytes_feature(tf.compat.as_bytes(uds.tostring()))}
    # Create an example protocol buffer
    example = tf.train.Example(features=tf.train.Features(feature=feature))
    # Serialize to string and write on the file
    writer.write(example.SerializeToString())
writer.close()
sys.stdout.flush()

# open the TFRecords file
test_filename = 'test.tfrecords'  # address to save the TFRecords file
writer = tf.python_io.TFRecordWriter(test_filename)
for i in range(len(test_addrs)):
    # print how many images are saved every 1000 images
    if not i % 1000:
        print ('Test data: {}/{}'.format(i, len(test_addrs)))
        sys.stdout.flush()
    # Load the image
    photo = load_image(test_addrs[i][0])
    uds = load_image(test_addrs[i][1])
    label = test_labels[i]
    # Create a feature
    feature = {'label': _int64_feature(label),
               'photo': _bytes_feature(tf.compat.as_bytes(photo.tostring())),
               'uds': _bytes_feature(tf.compat.as_bytes(uds.tostring()))}
    # Create an example protocol buffer
    example = tf.train.Example(features=tf.train.Features(feature=feature))
    # Serialize to string and write on the file
    writer.write(example.SerializeToString())
writer.close()
sys.stdout.flush()

Train data: 0/6324
Train data: 1000/6324
Train data: 2000/6324
Train data: 3000/6324
Train data: 4000/6324
Train data: 5000/6324
Train data: 6000/6324
Val data: 0/2108
Val data: 1000/2108
Val data: 2000/2108
Test data: 0/2108
Test data: 1000/2108
Test data: 2000/2108