Baza MNIST, w jaki sposób załadować interesujące mnie dane?

Cześć, chcę na bazie MNIST wytrenować sieć neuronową, jednak chciałbym, aby potrafiła ona rozpoznawać jedynie cyfrę 1 i 2. Dlatego też w celu jej nauki chciałbym załadować jej jedynie obrazy z cyframi 1 i 2. W tej chwili przy użyciu Tensorflow dane ładuję:

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

Co za problem, przecież Możesz sobie ściągnąć MNIST ręcznie z netu, np. stąd: http://yann.lecun.com/exdb/mnist/ i teraz dzielić dowolnie.