mnist数据集保存为图片

#coding: utf-8
from tensorflow.examples.tutorials.mnist import input_data
import scipy.misc
import os
import numpy as np

# 读取MNIST数据集。如果不存在会事先下载。
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

# 我们把原始图片保存在MNIST_data/raw/文件夹下
# 如果没有这个文件夹会自动创建
train_set_save_dir = 'MNIST_data/raw/train/'
test_set_save_dir = 'MNIST_data/raw/test/' 

if os.path.exists(train_set_save_dir) is False:
    os.makedirs(train_set_save_dir)

if os.path.exists(test_set_save_dir) is False:
    os.makedirs(test_set_save_dir)

n_classes = 10
train_set_dirs = []
test_set_dirs = []
train_file = []
test_file = []
    
for i in range(0, n_classes):
    train_dir = train_set_save_dir + str(i) + "/"
    train_set_dirs.append(train_dir)
    test_dir = test_set_save_dir + str(i) + "/"
    test_set_dirs.append(test_dir)
    
    train_file.append([])
    test_file.append([])
    
    if os.path.exists(train_dir) is False:
        os.makedirs(train_dir)
    if os.path.exists(test_dir) is False:
        os.makedirs(test_dir)

print(mnist.train.images.shape)
print(mnist.test.images.shape)
print(mnist.train.labels.shape)
print(mnist.test.labels.shape)

print(mnist.train.labels[0:10])
print(np.argmax(mnist.train.labels[0:10,:], axis=1))

trainImageCounts = mnist.train.images.shape[0]
testImageCounts = mnist.test.images.shape[0]

#exit()

# 保存训练集图片
for i in range(0, trainImageCounts):
    # 
    classIndex = np.argmax(mnist.train.labels[i], axis=0)
    
   
    train_file[classIndex].append('mnist_train_%d.jpg' % i)
    #continue
    
    # 保存文件的格式为 mnist_train_0.jpg, mnist_train_1.jpg, ... ,mnist_train_19.jpg
    #filename = train_set_save_dir + 'mnist_train_%d.jpg' % i
    filename = train_set_dirs[classIndex] + 'mnist_train_%d.jpg' % i
    
    # 请注意，mnist.train.images[i, :]就表示第i张图片（序号从0开始）
    image_array = mnist.train.images[i, :]
    
    # TensorFlow中的MNIST图片是一个784维的向量，我们重新把它还原为28x28维的图像。
    image_array = image_array.reshape(28, 28)
    
    
    # 将image_array保存为图片
    # 先用scipy.misc.toimage转换为图像，再调用save直接保存。
    scipy.misc.toimage(image_array, cmin=0.0, cmax=1.0).save(filename)
    
    
# 保存测试集图片
for i in range(0, testImageCounts):
    # 
    classIndex = np.argmax(mnist.test.labels[i], axis=0)
    
    
    test_file[classIndex].append('mnist_test_%d.jpg' % i)
    #continue
    
    # 保存文件的格式为 mnist_train_0.jpg, mnist_train_1.jpg, ... ,mnist_train_19.jpg
    #filename = test_set_save_dir + 'mnist_test_%d.jpg' % i
    filename = test_set_dirs[classIndex] + 'mnist_test_%d.jpg' % i
  
    # 请注意，mnist.test.images[i, :]就表示第i张图片（序号从0开始）
    image_array = mnist.test.images[i, :]
    
    # TensorFlow中的MNIST图片是一个784维的向量，我们重新把它还原为28x28维的图像。
    image_array = image_array.reshape(28, 28)
    
    # 将image_array保存为图片
    # 先用scipy.misc.toimage转换为图像，再调用save直接保存。
    scipy.misc.toimage(image_array, cmin=0.0, cmax=1.0).save(filename)

for i in range(0, n_classes):
    trainLabelFile = open(train_set_dirs[i] + "list.txt", "w")
    for picFlie in train_file[i]:
        trainLabelFile.write(picFlie + "\n")
    trainLabelFile.close()

    testLabelFile = open(test_set_dirs[i] + "list.txt", "w")
    for picFlie in test_file[i]:
        testLabelFile.write(picFlie + "\n")
    testLabelFile.close()

print('Please check: %s ' % train_set_save_dir)
print('Please check: %s ' % test_set_save_dir)