How to remove certain classes (labels and images) from existing TensorFlow datasets? (Fashion MNIST)

Viewed 1130

I am just getting started with learning more about TensorFlow and numpy. I am currently using the Fashion MNIST dataset with TensorFlow which includes 10 types of clothing. However, I want to be able to edit the numpy arrays that contain these datasets to remove all images and labels that are not 'T-shirts' and 'Shirts' and 'Trousers'. Essentially I just want to create a dataset from Fashion MNIST with only those 3 types.

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

Above is how I am importing the dataset currently as so far I am understanding that there are a few different ways to import the dataset before preprocessing it. How do I make sure I correctly remove the labels and its corresponding pictures so that the resulting labels and images still correspond with each other?

1 Answers
from tensorflow.keras.datasets import fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

# sorting based on index
idx = np.argsort(train_labels)
train_images = train_images[idx]
train_labels = train_labels[idx]

idx = np.argsort(test_labels)
test_images = test_images[idx]
test_labels = test_labels[idx]

labels = ["T-Shirt", "Trouser", "Pullover", "Dress", "Coat", 
          "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]

label_mapping = dict(zip(labels, range(10)))

def get_data(mapping, classes):
    X_train, X_test, y_train, y_test = [], [], [], []
    for cls in classes:
        idx = mapping[cls]
        start = idx*6000
        end = idx*6000+6000
        X_train.append(train_images[start: end])
        y_train.append(train_labels[start: end])
        start = idx*1000
        end = idx*1000+1000
        X_test.append(test_images[start: end])
        y_test.append(test_labels[start: end])
    return X_train, X_test, y_train, y_test


X_train, X_test, y_train, y_test = get_data(label_mapping, 
                                            classes=["T-Shirt", "Shirt", "Trouser"])

You can find the mapping between classes and their labels here

Related