import tensorflow as tf
from PIL import Image
import os
import numpy as np
from sklearn.model_selection import train_test_split
import time
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
# 读取所有图片 统一尺寸 转为像素矩阵
VGG_MEAN = [103.939, 116.778, 123.68]
dir_path = 'newfile'
path_Vgg16 = r'C:\Users\machenike\Desktop\vgg16.npy'
# 所有样本
data_all = []
# 所有标签
labels_all = []
for sub_dir in os.listdir(dir_path):
# print(sub_dir)
i = 0
for img_path in os.listdir(os.path.join(dir_path, sub_dir)):
# 图片全路径
imgs = Image.open(os.path.join(dir_path, sub_dir, img_path))
# 转换通道
imgs = imgs.convert('RGB')
# 重置尺寸
imgs = imgs.resize((224, 224))
img_array = np.array(imgs)
data_all.append(img_array)
# print(img_array.shape)
# 制作标签
label = str.split(img_path, '.')[0]
labels_all.append(label)
# 转为数组
data_all = np.array(data_all)
labels_all = np.asarray(labels_all)
labels_all = labels_all.astype(np.int)
print(data_all.shape)
print(labels_all.shape)
# 数据集切分和打乱
x_train, x_test, y_train, y_test = train_test_split(data_all, labels_all, test_size=0.1, shuffle=20)
# 载入模型
data_dict = np.load(path_Vgg16, allow_pickle=True, encoding='latin1').item()
class VGGNet:
def __init__(self, data_dict):
self.data_dict = data_dict
# 获取卷积层卷积核
def get_conv_filter(self, name):
return tf.constant(self.data_dict[name][0], name='conv')
# 获取全连接层权重
def get_fc_weight(self, name):
return tf.constant(self.data_dict[name][0], name='fc')
# 获取偏置
def get_bias(self, name):
return tf.constant(self.data_dict[name][1], name='bias')
# 构建卷积层
def conv_layer(self, x, name):
with tf.name_scope(name):
conv_w = self.get_conv_filter(name)
conv_b = self.get_bias(name)
h = tf.nn.conv2d(x, conv_w, [1, 1, 1, 1], padding='SAME')
h = tf.nn.bias_add(h, conv_b)
h = tf.nn.relu(h)
return h
# 池化层
def pooling_layer(self, x, name):
return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)
# 构建全连接
def fc_layer(self, x, name, activation=tf.nn.relu):
with tf.name_scope(name):
fc_w = self.get_fc_weight(name)
fc_b = self.get_bias(name)
h = tf.matmul(x, fc_w)
h = tf.nn.bias_add(h, fc_b)
if activation is None:
return h
else:
return activation(h)
# 构建展平层
def flatten_layer(self, x, name):
with tf.name_scope(name):
x_shape = x.get_shape().as_list()
dim = 1
for d in x_shape:
dim *= d
x = tf.reshape(x, [-1, dim])
return x
def build(self, x_rgb):
r, g, b = tf.split(x_rgb, [1, 1, 1], axis=3)
x_bgr = tf.concat([b - VGG_MEAN[0],
g - VGG_MEAN[1],
r - VGG_MEAN[2]],
axis=3)
assert x_bgr.get_shape().as_list()[1:] == [224, 224, 3]
self.conv1_1 = self.conv_layer(x_bgr, 'conv1_1')
self.conv1_2 = self.conv_layer(self.conv1_1, 'conv1_2')
self.pool1 = self.pooling_layer(self.conv1_2, 'pool1')
self.conv2_1 = self.conv_layer(self.pool1, 'conv2_1')
self.conv2_2 = self.conv_layer(self.conv2_1, 'conv2_2')
self.pool2 = self.pooling_layer(self.conv2_2, 'pool2')
self.conv3_1 = self.conv_layer(self.pool2, 'conv3_1')
self.conv3_2 = self.conv_layer(self.conv3_1, 'conv3_2')
self.conv3_3 = self.conv_layer(self.conv3_2, 'conv3_3')
self.pool3 = self.pooling_layer(self.conv3_3, 'pool3')
self.conv4_1 = self.conv_layer(self.pool3, 'conv4_1')
self.conv4_2 = self.conv_layer(self.conv4_1, 'conv4_2')
self.conv4_3 = self.conv_layer(self.conv4_2, 'conv4_3')
self.pool4 = self.pooling_layer(self.conv4_3, 'pool4')
self.conv5_1 = self.conv_layer(self.pool4, 'conv5_1')
self.conv5_2 = self.conv_layer(self.conv5_1, 'conv5_2')
self.conv5_3 = self.conv_layer(self.conv5_2, 'conv5_3')
self.pool5 = self.pooling_layer(self.conv5_3, 'pool5')
return self.pool5
# self.flatten5 = self.flatten_layer(self.pool5, 'flatten')
# self.fc6 = self.fc_layer(self.flatten5, 'fc6')
# self.fc7 = self.fc_layer(self.fc6, 'fc7')
# self.fc8 = self.fc_layer(self.fc7, 'fc8', activation=None)
# self.prob = tf.nn.softmax(self.fc8, name='prob')
x = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
y = tf.placeholder(tf.int64, shape=[None])
x_model = VGGNet(data_dict)
pool5 = x_model.build(x)
flatten = tf.layers.flatten(pool5)
# 构建全连接层
fc1 = tf.layers.dense(flatten, 128, activation=tf.nn.relu)
fc2 = tf.layers.dense(fc1, 128, activation=tf.nn.relu)
fc3 = tf.layers.dense(fc2, 64, activation=tf.nn.relu)
y_ = tf.layers.dense(fc3, 5)
# 损失函数
# loss = tf.losses.sparse_softmax_cross_entropy(logits=y_, labels=y)
# train
# train_op = tf.train.AdamOptimizer(0.001).minimize(loss)
# 预测
predict = tf.argmax(y_, 1)
# 准确率
accuracy = tf.reduce_mean(tf.cast(tf.equal(predict, y), dtype=tf.float32))
# config = tf.ConfigProto()
# config.gpu_options.allow_growth = True
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# 读取保存好的模型 恢复之前的会话,不用做训练,直接测试
saver.restore(sess,save_path='./model/office_model.ckpt')
acc_all = []
for j in range(5):
x_test_batch, y_test_batch = x_test[j * 20: (j + 1) * 20, :, :, :], y_test[j * 20: (j + 1) * 20]
acc_val = sess.run(accuracy, feed_dict={x: x_test_batch, y: y_test_batch})
acc_all.append(acc_val)
print('测试集准确率:', sess.run(tf.reduce_mean(acc_all)))
'''
测试集准确率: 0.91999996
'''
提取利用迁移学习保存好的模型应用到工业界
最新推荐文章于 2024-04-15 17:11:08 发布