Python3.7的Tensorflow2.0 图像分类的模型生成与读取

源码下载

链接: https://pan.baidu.com/s/1mR8e8lBJV_69sqNt_Irmvw 提取码: n553

首先是读取照片，照片放在dataset文件夹内。每种是一个文件夹。
 

 

 

新建一个py文件，utils_paths.py。读取照片。

import os
 
 
image_types = (".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff")
 
 
def list_images(basePath, contains=None):
    # 返回有效的图片路径数据集
    return list_files(basePath, validExts=image_types, contains=contains)
 
 
def list_files(basePath, validExts=None, contains=None):
    # 遍历图片数据目录，生成每张图片的路径
    for (rootDir, dirNames, filenames) in os.walk(basePath):
        # 循环遍历当前目录中的文件名
        for filename in filenames:
            # if the contains string is not none and the filename does not contain
            # the supplied string, then ignore the file
            if contains is not None and filename.find(contains) == -1:
                continue
 
            # 通过确定.的位置，从而确定当前文件的文件扩展名
            ext = filename[filename.rfind("."):].lower()
 
            # 检查文件是否为图像，是否应进行处理
            if validExts is None or ext.endswith(validExts):
                # 构造图像路径
                imagePath = os.path.join(rootDir, filename)
                yield imagePath

然后是图片用keras训练，生成模型，尾缀是.h5。然后可以生成frozen的pb文件。读取pb文件，都有。

# 导入所需工具包
#from CNN_net import SimpleVGGNet
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
import utils_paths
import matplotlib.pyplot as plt
import numpy as np
import argparse
import random
import pickle
import cv2
import os
from keras.models import Sequential
from keras.layers.normalization import BatchNormalization
from keras.layers.convolutional import Conv2D
from keras.layers.convolutional import MaxPooling2D
from keras.initializers import TruncatedNormal
from keras.layers.core import Activation
from keras.layers.core import Flatten
from keras.layers.core import Dropout
from keras.layers.core import Dense
from keras import backend as K
from keras.models import load_model
from keras.optimizers import Adam
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

def build(width, height, depth, classes):

		# initialize the model along with the input shape to be

		# "channels last" and the channels dimension itself

		model = Sequential()

		inputShape = (height, width, depth)

		chanDim = -1

 

		# if we are using "channels first", update the input shape

		# and channels dimension

		if K.image_data_format() == "channels_first":

			inputShape = (depth, height, width)

			chanDim = 1

# CONV => RELU => POOL

		model.add(Conv2D(32, (3, 3), padding="same",

			input_shape=inputShape))

		model.add(Activation("relu"))

		model.add(BatchNormalization(axis=chanDim))

		model.add(MaxPooling2D(pool_size=(3, 3)))

		model.add(Dropout(0.25))

		# (CONV => RELU) * 2 => POOL

		model.add(Conv2D(64, (3, 3), padding="same"))

		model.add(Activation("relu"))

		model.add(BatchNormalization(axis=chanDim))

		model.add(Conv2D(64, (3, 3), padding="same"))

		model.add(Activation("relu"))

		model.add(BatchNormalization(axis=chanDim))

		model.add(MaxPooling2D(pool_size=(2, 2)))

		model.add(Dropout(0.25))

		# (CONV => RELU) * 2 => POOL

		model.add(Conv2D(128, (3, 3), padding="same"))

		model.add(Activation("relu"))

		model.add(BatchNormalization(axis=chanDim))

		model.add(Conv2D(128, (3, 3), padding="same"))

		model.add(Activation("relu"))

		model.add(BatchNormalization(axis=chanDim))

		model.add(MaxPooling2D(pool_size=(2, 2)))

		model.add(Dropout(0.25))

		# first (and only) set of FC => RELU layers

		model.add(Flatten())

		model.add(Dense(1024))

		model.add(Activation("relu"))

		model.add(BatchNormalization())

		model.add(Dropout(0.5))

 

		# softmax classifier

		model.add(Dense(classes))

		model.add(Activation("softmax"))


		# return the constructed network architecture

		return model

# 读取数据和标签
print("------开始读取数据------")
data = []
labels = []
EPOCHS = 20
BS =128
# 拿到图像数据路径，方便后续读取
imagePaths = sorted(list(utils_paths.list_images('./dataset')))
random.seed(42)
random.shuffle(imagePaths)

# 遍历读取数据
for imagePath in imagePaths:
    # 读取图像数据
    image = cv2.imread(imagePath)
    image = cv2.resize(image, (64, 64))
    data.append(image)
    # 读取标签
    label = imagePath.split(os.path.sep)[-2]
    labels.append(label)

# 对图像数据做scale操作
data = np.array(data, dtype="float32") / 255.0
labels = np.array(labels)

# 数据集切分
(trainX, testX, trainY, testY) = train_test_split(data,labels, test_size=0.25, random_state=42)

# 转换标签为one-hot encoding格式
lb = LabelBinarizer()
trainY = lb.fit_transform(trainY)
testY = lb.transform(testY)

# 数据增强处理
aug = ImageDataGenerator(rotation_range=30, width_shift_range=0.1,
    height_shift_range=0.1, shear_range=0.2, zoom_range=0.2,
    horizontal_flip=True, fill_mode="nearest")
INIT_LR = 1e-3
# 建立卷积神经网络
model = build(width=64, height=64, depth=3,classes=len(lb.classes_))
opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)

# 设置初始化超参数



# 损失函数，编译模型
print("------准备训练网络------")

model.compile(loss="categorical_crossentropy", optimizer=opt,

	metrics=["accuracy"])

model.summary()
# 训练网络模型
H = model.fit_generator(aug.flow(trainX, trainY, batch_size=BS),validation_data=(testX, testY), steps_per_epoch=len(trainX) ,epochs=EPOCHS, verbose=1)




# 测试
print("------测试网络------")
predictions = model.predict(testX, batch_size=32)
print(classification_report(testY.argmax(axis=1),
    predictions.argmax(axis=1), target_names=lb.classes_))

# 绘制结果曲线
N = np.arange(0, EPOCHS)
plt.style.use("ggplot")
plt.figure()
plt.plot(N, H.history["loss"], label="train_loss")
plt.plot(N, H.history["val_loss"], label="val_loss")
plt.plot(N, H.history["accuracy"], label="train_acc")
plt.plot(N, H.history["val_accuracy"], label="val_acc")
plt.title("Training Loss and Accuracy")
plt.xlabel("Epoch #")
plt.ylabel("Loss/Accuracy")
plt.legend()
plt.savefig('./output/cnn_plot.png')

# 保存模型
print("------正在保存模型------")
model.save('./output/cnn.h5')
f = open('./output/cnn_lb.pickle', "wb")
f.write(pickle.dumps(lb))
f.close()


model=tf.keras.models.load_model('./output/cnn.h5')
#tf.keras.models.save_model(model,'./output/cnn.pb/') 
#tf.saved_model.save(model, "./output")

  # Convert Keras model to ConcreteFunction

full_model = tf.function(lambda x: model(x))

full_model = full_model.get_concrete_function(

tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))



    # Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()


layers = [op.name for op in frozen_func.graph.get_operations()]

print("-" * 50)

print("Frozen model layers: ")

for layer in layers:

        print(layer)



print("-" * 50)

print("Frozen model inputs: ")

print(frozen_func.inputs)

print("Frozen model outputs: ")

print(frozen_func.outputs)



    # Save frozen graph from frozen ConcreteFunction to hard drive

tf.io.write_graph(graph_or_graph_def=frozen_func.graph,

                      logdir="./frozen_models",

                      name="frozen_graph.pb",

                      as_text=False)


#model=tf.keras.models.load_model('./frozen_models/frozen_graph.pb')
# 加载测试数据并进行相同预处理操作
image = cv2.imread('./pic/t3.jpg')
output = image.copy()
image = cv2.resize(image, (64, 64))

# scale图像数据
image = image.astype("float32") / 255.0

# 对图像进行拉平操作
image = image.reshape((1, image.shape[0], image.shape[1],image.shape[2]))

# 读取模型和标签
print("------读取模型和标签------")
#model = load_model('./output/cnn.model')
lb = pickle.loads(open('./output/cnn_lb.pickle', "rb").read())

# 预测
preds = model.predict(image)

# 得到预测结果以及其对应的标签
i = preds.argmax(axis=1)[0]
label = lb.classes_[i]

# 在图像中把结果画出来
text = "{}: {:.2f}%".format(label, preds[0][i] * 100)
cv2.putText(output, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7,(0, 0, 255), 2)

# 绘图
cv2.imshow("Image", output)
cv2.waitKey(0)


def wrap_frozen_graph(graph_def, inputs, outputs, print_graph=False):

    def _imports_graph_def():

        tf.compat.v1.import_graph_def(graph_def, name="")

 

    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])

    import_graph = wrapped_import.graph

 

    print("-" * 50)

    print("Frozen model layers: ")

    layers = [op.name for op in import_graph.get_operations()]

    if print_graph == True:

        for layer in layers:

            print(layer)

    print("-" * 50)

 

    return wrapped_import.prune(

        tf.nest.map_structure(import_graph.as_graph_element, inputs),

        tf.nest.map_structure(import_graph.as_graph_element, outputs))

 

 




with tf.io.gfile.GFile("./frozen_models/frozen_graph.pb", "rb") as f:        
	graph_def = tf.compat.v1.GraphDef()        
	loaded = graph_def.ParseFromString(f.read())
frozen_func = wrap_frozen_graph(graph_def=graph_def,  
   inputs=["x:0"],                                    
   outputs=["Identity:0"], 
   print_graph=True)

print("-" * 50)   
print("Frozen model inputs: ")   
print(frozen_func.inputs)    
print("Frozen model outputs: ")    
print(frozen_func.outputs)     
# Get predictions for test images    
predictions = frozen_func(x=tf.constant(image))    
# # Print the prediction for the first image    
print("-" * 50)    
print("Example prediction reference:")    
print(predictions[0].numpy())