Unable to read a TFRecord file as an Image

I have a lot of TFRecord images from google earth engine that I am trying to read in as images/maps. I tried to follow this page:

and my code is as follows:

raw_image_dataset = tf.data.TFRecordDataset('uganda_2005_01.tfrecord')

# Create a dictionary describing the features.
image_feature_description = {
    'height': tf.io.FixedLenFeature([], tf.int64),
    'width': tf.io.FixedLenFeature([], tf.int64),
    'depth': tf.io.FixedLenFeature([], tf.int64),
    'label': tf.io.FixedLenFeature([], tf.int64),
    'image_raw': tf.io.FixedLenFeature([], tf.string),
}

def _parse_image_function(example_proto):
  # Parse the input tf.train.Example proto using the dictionary above.
  return tf.io.parse_single_example(example_proto, image_feature_description)

parsed_image_dataset = raw_image_dataset.map(_parse_image_function)
parsed_image_dataset

for image_features in parsed_image_dataset:
  image_raw = image_features['image_raw'].numpy()
  display.display(display.Image(data=image_raw))

I keep getting this error that I am not sure how to fix:


InvalidArgumentError Traceback (most recent call last)
/var/folders/my/mlj39v8d0jl1gjcg98xyly3h0000gn/T/ipykernel_99279/1152503565.py in
----> 1 for image_features in parsed_image_dataset:
2 image_raw = image_features[‘image_raw’].numpy()
3 display.display(display.Image(data=image_raw))

~/opt/miniconda3/envs/CNNRep/lib/python3.7/site-packages/tensorflow_core/python/data/ops/iterator_ops.py in next(self)
620
621 def next(self): # For Python 3 compatibility
→ 622 return self.next()
623
624 def _next_internal(self):

~/opt/miniconda3/envs/CNNRep/lib/python3.7/site-packages/tensorflow_core/python/data/ops/iterator_ops.py in next(self)
664 “”“Returns a nested structure of Tensors containing the next element.”""
665 try:
→ 666 return self._next_internal()
667 except errors.OutOfRangeError:
668 raise StopIteration

~/opt/miniconda3/envs/CNNRep/lib/python3.7/site-packages/tensorflow_core/python/data/ops/iterator_ops.py in _next_internal(self)
649 self._iterator_resource,
650 output_types=self._flat_output_types,
→ 651 output_shapes=self._flat_output_shapes)
652
653 try:

~/opt/miniconda3/envs/CNNRep/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_dataset_ops.py in iterator_get_next_sync(iterator, output_types, output_shapes, name)
2671 else:
2672 message = e.message
→ 2673 _six.raise_from(_core._status_to_exception(e.code, message), None)
2674 # Add nodes to the TensorFlow graph.
2675 if not isinstance(output_types, (list, tuple)):

~/opt/miniconda3/envs/CNNRep/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: {{function_node __inference_Dataset_map__parse_image_function_61}} Feature: depth (data type: int64) is required but could not be found.
[[{{node ParseSingleExample/ParseSingleExample}}]] [Op:IteratorGetNextSync]

Any help would be greatly appreciated.