JGuizard
9/25/2019 - 12:40 PM

Extract images related to just some classes from MNIST dataset

Extract images related to just some classes from MNIST dataset

from PIL import Image
from os import mkdir

BASE_PATH = "mnist/"

classes = [0,1]
counter = [1,1]
limit = 1000

for c in classes:
    mkdir(BASE_PATH+str(c))

for x, y in zip(X_train, y_train):
    if(y in classes):
        if(counter[y]<=limit):
            im = Image.fromarray(x.reshape(28,28))
            im.save(BASE_PATH+"%d/%d.jpg" % (y, counter[y]))
            counter[y]+=1
        if(counter[0]>1000 and counter[1]>1000):
            break