如何 select 在 cifar-10 中具体 class

Question

I would like to know how to select the specific class in cifar-10. For example, I want 7, "horse" class in cifar-10. And I wrote the below code. But the obtained data is not what I want because it's wrong shape.

Please enlighten me on the specifics.

from keras.datasets import cifar10

(X_train, Y_train), (X_test, Y_test) = cifar10.load_data()

print('X_train shape: {0}, Y_train shape: {1}'.format(X_train.shape, Y_train.shape))

X_train shape: (50000, 32, 32, 3), Y_train shape: (50000, 1)

The below code is wrong.

import numpy as np

filter = np.where(Y_train == 7)

X_train = X_train[filter]
Y_train = Y_train[filter]

print('X_train shape: {0}, Y_train shape: {1}'.format(X_train.shape, Y_train.shape))

X_train shape: (5000, 32, 3), Y_train shape: (5000,)

The expected output is below

X_train shape: (5000, 32, 32, 3), Y_train shape: (5000,)

Answer 1

For the slicing, do something like:

X_train = X_train[filter[0], ...]
Y_train = Y_train[filter[0], ...]

And the shapes would be

X_train shape: (5000, 32, 32, 3), Y_train shape: (5000, 1)