caffe_mnist_image_extractor.py

Go to the documentation of this file.

#!/usr/bin/env python
"""Extracts mnist image data from the Caffe data files and stores them in numpy arrays
Usage
    python caffe_mnist_image_extractor.py -d path_to_caffe_data_directory -o desired_output_path
 
Saves the first 10 images extracted as input10.npy, the first 100 images as input100.npy, and the
corresponding labels to labels100.txt.
 
Tested with Caffe 1.0 on Python 2.7
"""
import argparse
import os
import struct
import numpy as np
from array import array
 
 
if __name__ == "__main__":
    # Parse arguments
    parser = argparse.ArgumentParser('Extract Caffe mnist image data')
    parser.add_argument('-d', dest='dataDir', type=str, required=True, help='Path to Caffe data directory')
    parser.add_argument('-o', dest='outDir', type=str, default='.', help='Output directory (default = current directory)')
    args = parser.parse_args()
 
    images_filename = os.path.join(args.dataDir, 'mnist/t10k-images-idx3-ubyte')
    labels_filename = os.path.join(args.dataDir, 'mnist/t10k-labels-idx1-ubyte')
 
    images_file = open(images_filename, 'rb')
    labels_file = open(labels_filename, 'rb')
    images_magic, images_size, rows, cols = struct.unpack('>IIII', images_file.read(16))
    labels_magic, labels_size = struct.unpack('>II', labels_file.read(8))
    images = array('B', images_file.read())
    labels = array('b', labels_file.read())
 
    input10_path   = os.path.join(args.outDir, 'input10.npy')
    input100_path  = os.path.join(args.outDir, 'input100.npy')
    labels100_path = os.path.join(args.outDir, 'labels100.npy')
 
    outputs_10  = np.zeros(( 10, 28, 28, 1), dtype=np.float32)
    outputs_100 = np.zeros((100, 28, 28, 1), dtype=np.float32)
    labels_output = open(labels100_path, 'w')
    for i in xrange(100):
        image = np.array(images[i * rows * cols : (i + 1) * rows * cols]).reshape((rows, cols)) / 256.0
        outputs_100[i, :, :, 0] = image
 
        if i < 10:
            outputs_10[i, :, :, 0] = image
 
        if i == 10:
            np.save(input10_path, np.transpose(outputs_10, (0, 3, 1, 2)))
            print "Wrote", input10_path
 
        labels_output.write(str(labels[i]) + '\n')
 
    labels_output.close()
    print "Wrote", labels100_path
 
    np.save(input100_path, np.transpose(outputs_100, (0, 3, 1, 2)))
    print "Wrote", input100_path