Write .nii data to a TFRecords database

write_TFRecord.py

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

def _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

# path to save the TFRecords file
train_filename = 'train.tfrecords'

# open the file
writer = tf.python_io.TFRecordWriter(train_filename)

# iterate through all .nii files:
for meta_data in all_filenames:

    # Load the image and label
    img, label = load_img(meta_data, reader_params)
    
    # Create a feature
    feature = {'train/label': _int64_feature(label),
               'train/image': _float_feature(img.ravel())}
               
    # 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()