import torchvision
from PIL import Image

data = torchvision.datasets.MNIST('data/', train=True, download=True)

def new_image():
    width = 28 * 100
    height = 28 * 70
    return Image.new('L', (width, height))
    

images = [new_image() for i in range(10)]
counts = [0 for _ in range(10)]

for i in range(len(data)):
    img, label = data[i]
    images[label].paste(img, (28 * (counts[label] % 100), 28 * (counts[label] // 100)))
    counts[label] += 1

for i in range(10):
    images[i].save('data/images-{}.png'.format(i))