COCO is available also in TF datasets I suppose you could filter like in:

  

      github.com/tensorflow/datasets
  

  
    
  
	  
  

  
    
      How to filter the dataset to get images from a specific class? 1
    

    
      
        opened 
        
          
        
        Apr 22, 2020
      
      

        
          closed 
        
          
        
        Apr 23, 2020
      
        

      
        
          
          rao208
        
      
    

    
        
          help
        
    
  


  
    **What I need help with / What I was wondering**

I want to use pre-trained ne…tworks to train my dataset. I am a novice at this and right now just exploring and playing with the basic dataset cats_vs_dogs as given in the link:

[https://www.tensorflow.org/tutorials/images/transfer_learning ](URL)

I am using the exact same code with some modifications like one-hot encoding and the architecture of my own.

I was wondering how to plot the images of only one class in cat_vs_dog dataset i.e. either cat or dog. In other words, I want to separate the dataset belonging to class 'cat' from that belonging to class 'dog'. I would like to do this not for training purpose, but to achieve some different goals

**What I've tried so far**

There is an answer on StackOverflow ([https://stackoverflow.com/questions/55731774/filter-dataset-to-get-just-images-from-specific-class](URL))

Unfortunately, this didn't give the desired output because of the following issues:

* **Issue 1:** label = x['label'] in predicate function gave me error

> TypeError: Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got 'label'

So, I replaced x['label'] with x[0] i.e. class 'cat'. The predicate function looks something like this:

```
def predicate(x, allowed_labels=tf.constant([0.])):
    label = x[0]
    isallowed = tf.equal( tf.cast(allowed_labels, tf.float32), tf.cast(label, tf.float32))
    reduced = tf.reduce_sum(tf.cast(isallowed, tf.float32))
    return tf.greater(reduced, tf.constant(0.))

train = raw_train.map(format_example)
validation = raw_validation.map(format_example)
test = raw_test.map(format_example)

train = train.filter(predicate)
validation = validation.filter(predicate)
test = test.filter(predicate)

for image, label in train.take(10):
    print("label---", label)
    plt.figure()
    plt.imshow(image)
    plt.title(get_label_name(label))

plt.show()
```

The output is:

```
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
label--- tf.Tensor(1, shape=(), dtype=int64)
```
So, it is basically giving me only class 'dog', even when I input class 'cat'. Maybe this is because predicate function returns **tf.greater(reduced, tf.constant(0.))**.

* **Issue 2:**

In a naive attempt, to overcome Issue 1, I changed return to 
`return tf.greater_equal(reduced, tf.constant(0.))` and `return tf.greater(reduced, tf.constant(-1.))`

and due to obvious reasons, it didn't work. 

* **Issue 3:**

Also, let's say I have more than two classes like cifar10 dataset and I would like to keep class 4 and 7 and filter out the rest of the classes. Then, how should I modify the code? As mentioned before x['label'] is giving me an error.

Is it possible to make predicate function more generic, so that I can keep N number of classes and filter out the rest of the classes? or is there any other way to filter the dataset to get images from a specific class?

**Environment information**
* Operating System: <windows>
* Distribution: Anaconda
* Python version: <3.7.7>
* Tensorflow 2.1
* tensorflow_datasets 1.2.0

Please let me know if any more information is required. Thank you for your help in advance.

Best Regards