Function main

samples/python/scripts/download_mnist_pgms.py:46–65

()

    return list(raw_buf[8:].astype(np.int32).reshape(num_labels))

def main():
    parser = argparse.ArgumentParser(description="Extracts 10 PGM files from the MNIST dataset", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-o", "--output", help="Path to the output directory.", default=os.getcwd())

    args, _ = parser.parse_known_args()

    with urllib.request.urlopen("http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz") as res:
        data = load_mnist_data(gzip.decompress(res.read()))

    with urllib.request.urlopen("http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz") as res:
        labels = load_mnist_labels(gzip.decompress(res.read()))

    output_dir = args.output

    # Find one image for each digit.
    for i in range(10):
        index = labels.index(i)
        image = Image.fromarray(data[index], mode="L")
        path = os.path.join(output_dir, "{:}.pgm".format(i))
        image.save(path)

if __name__ == '__main__':
    main()