Чтение изображений из TFrecord с использованием API набора данных и отображение их в блокноте Jupyter

Я создал tfrecord из папки изображений, теперь я хочу перебрать записи в файле TFrecord с помощью API набора данных и показать их в блокноте Jupyter. Однако у меня возникают проблемы с чтением файла tfrecord.

Код, который я использовал для создания TFRecord

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

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

def generate_tfr(image_list):
    with tf.python_io.TFRecordWriter(output_path) as writer:
        for image in images:
            image_bytes = open(image,'rb').read()
            image_array = imread(image)
            image_shape = image_array.shape
            image_x, image_y, image_z = image_shape[0],image_shape[1], image_shape[2]

            data = {

              'image/bytes':_bytes_feature(image_bytes),
              'image/x':_int64_feature(image_x),
              'image/y':_int64_feature(image_y),
              'image/z':_int64_feature(image_z)
            }

            features = tf.train.Features(feature=data)
            example = tf.train.Example(features=features)
            serialized = example.SerializeToString()
            writer.write(serialized)

Код для чтения TFRecord

#This code is incomplete and has many flaws. 
#Please give some suggestions in correcting this code if you can

def parse(serialized):
    features = \
    {
        'image/bytes': tf.FixedLenFeature([], tf.string),
        'image/x': tf.FixedLenFeature([], tf.int64),
        'image/y': tf.FixedLenFeature([], tf.int64),
        'image/z': tf.FixedLenFeature([], tf.int64)
    }

    parsed_example = tf.parse_single_example(serialized=serialized,features=features)
    image = parsed_example['image/bytes']
    image = tf.decode_raw(image,tf.uint8)
    x = parsed_example['image/x'] # breadth
    y = parsed_example['image/y'] # height
    z = parsed_example['image/z'] # depth
    image = tf.cast(image,tf.float32)

    # how can I reshape image tensor here? tf.reshape throwing some weird errors.

    return {'image':image,'x':x,'y':y,'z':z}



dataset = tf.data.TFRecordDataset([output_path])
dataset.map(parse)
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()
epoch = 1

with tf.Session() as sess:
    for _ in range(epoch):
    img = next_element.eval()
    print(img)
    # when I print image, it shows byte code.
    # How can I convert it to numpy array and then show image on my jupyter notebook ?

Я никогда раньше не работал ни с чем из этого, и я застрял в чтении TFRecords. Пожалуйста, ответьте, как перебрать содержимое TFrecords и показать их в блокноте Jupyter. Не стесняйтесь исправлять/оптимизировать обе части кода. Это бы мне очень помогло.