用几行代码将MNIST数据集图片提取并分类存放

#inputs date
import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import tensorflow.examples.tutorials.mnist.input_data as input_data
mnist_path="D:\\log\\MNIST"
save_path="D:\\log\\MNIST_JPG"
mnist=input_data.read_data_sets(mnist_path,one_hot=True)
def decode(images,labels,path):
    for index,vec in enumerate(images):
        image=vec.reshape([28,28])    #遍历图片
        label=np.argmax(labels[index])#图片对应标签
        classes_path=os.path.join(path,str(label)) #图片分类到每个文件夹,得到文件夹地址
        if not os.path.exists(classes_path):os.makedirs(classes_path) #不存在就创建
        image_path=os.path.join(classes_path,str(index)+".jpg") #图片地址
        matplotlib.image.imsave(image_path,image,cmap="binary")#保存图片
decode(mnist.validation.images,mnist.validation.labels,os.path.join(save_path,"validation"))
print("finish")
decode(mnist.train.images,mnist.train.labels,os.path.join(save_path,"train"))
print("finish")
decode(mnist.test.images,mnist.test.labels,os.path.join(save_path,"test"))
print("finish")

UPD:
对于较高版本的tf这个方法可能并不行
会出现找不到input_data的问题
原因是高版本已经废弃了这个接口
需要用到tf.keras的接口

import os
import numpy as np
import matplotlib
import tensorflow as tf
save_path="D:\\log\\MNIST_PNG"
(x_train,y_train),(x_test,y_test)=tf.keras.datasets.mnist.load_data()
def decode(images,labels,path):
	print("start")
        image=vec.reshape([28,28])    #遍历图片
        label=labels[index]#图片对应标签
        classes_path=os.path.join(path,str(label)) #图片分类到每个文件夹,得到文件夹地址
        if not os.path.exists(classes_path):os.makedirs(classes_path) #不存在就创建
        image_path=os.path.join(classes_path,str(index)+".png") #图片地址
        matplotlib.image.imsave(image_path,image,cmap="binary")#保存图片    
decode(x_train,y_train,os.path.join(save_path,"train"))
print("finish")
decode(x_test,y_test,os.path.join(save_path,"test"))
print("finish")

结果: