Filter Tensorflow dataset by Class

user_6106508 0
Viewed 406

I would like to filter a tensorflow dataset to output only values of a specific class/label. How would we do that for the code below?

Thanks,

image_feature_description = {
    'label': tf.io.FixedLenFeature([], tf.string),
    'image': tf.io.FixedLenFeature([100, 100, 3], tf.float32),
}

def parse_tfrecord(example_proto):
  features = tf.io.parse_example(example_proto, image_feature_description)
  label = features['label']
  image = features['image']
  return image

  dataset = dataset.map(parse_tfrecord).batch(batch_size)
tensorflow-datasets tensorflow python
1 Answers

modify parse_tfrecord function to return both label and image as following:

def parse_tfrecord(example_proto):
  ... # parsing an example
  return image, label

then, add a filter which keeps label==MyLabel only between map and batch op:

dataset = dataset.map(parse_tfrecord) \
  .filter(lambda image, label: label == MY_LABEL) \
  .map(lambda image, label: image) \  # add this if you want image only
  .batch(...)
user_15618740
Related