Skip to content

Instantly share code, notes, and snippets.

@lartpang
Last active July 9, 2022 04:26
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save lartpang/e2343fb06f25b36a8af03e68923dea06 to your computer and use it in GitHub Desktop.
Save lartpang/e2343fb06f25b36a8af03e68923dea06 to your computer and use it in GitHub Desktop.
Keras预训练模型/按需分配显存/冻结层/独热编码解码输出/写文件/IoU/导出xml
#!/usr/bin/env python
import os
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
import numpy as np
from keras.layers import Dense, GlobalAveragePooling2D
from keras.models import Model
from keras.optimizers import Adam
from keras.preprocessing import image
from keras.applications.vgg16 import VGG16
config = tf.ConfigProto()
config.gpu_options.allow_growth=True #不全部占满显存, 按需分配
sess = tf.Session(config=config)
KTF.set_session(sess)
# 导入训练数据
x_train = np.load('x_train.npy')
y_train = np.load('y_train.npy')
# 搭建模型
base_model = VGG16(weights='imagenet', include_top=False)
adam = Adam(lr=0.0001, beta_1=0.99, beta_2=0.999, epsilon=1e-08)
# 添加新层
old_layer = base_model.output
new_layer = GlobalAveragePooling2D()(old_layer)
# 将四维转化为二维
new_layer = Dense(1024, activation='relu')(new_layer)
predict_layer = Dense(3, activation='softmax')(new_layer)
model = Model(input=base_model.input, output=predict_layer)
# 冻结层
for layer in base_model.layers:
layer.trainable = False
# 编译
model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=['accuracy'])
history = model.fit(x_train, y_train, batch_size=128, epochs=50,
verbose=2, validation_split=0.2)
# 测试数据
x_test = np.load('x_test.npy')
y_test = np.load('y_test.npy')
out = np.zeros((x_test.shape[0], 3))
out = model.predict(x_test)
y_test_pred = np.argmax(out, axis=1)
score = model.evaluate(x_test, y_test, verbose=2)
print("test score:", score[0])
print("test accuracy:", score[1])
# 导出数据
out_str = list(range(len(y_test_pred)))
for i, item in enumerate(y_test_pred):
if item == 0:
out_str[i] = 'A'
elif item == 1:
out_str[i] = 'B'
else:
out_str[i] = 'C'
out_str = ''.join(out_str)
with open("y_test_out.txt", 'w') as file_out:
file_out.write(out_str)
#!/usr/bin/env python
import os
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
import numpy as np
import keras
from keras.layers import Dense
from keras.models import Model
from keras.optimizers import Adam
from keras.preprocessing import image
from keras.applications.vgg16 import VGG16
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # 不全部占满显存, 按需分配
sess = tf.Session(config=config)
KTF.set_session(sess)
# 构建不带分类器的预训练模型
inputshape = (224, 224, 3)
base_model = VGG16(input_shape=inputshape,
weights='imagenet',
include_top=True)
old_layer = base_model.get_layer('fc2').output
predictions = Dense(4, activation='linear')(old_layer)
# 构建我们需要训练的完整模型
model = Model(inputs=base_model.input, outputs=predictions)
callbacks = [keras.callbacks.ModelCheckpoint(
'./model_data/weights_best_vgg16.hdf5',
verbose=1, save_weights_only=True, mode='max')]
optim = keras.optimizers.Adam(lr=1e-4, beta_1=0.99)
model.compile(optimizer=optim, loss='mse', metrics=['mae'])
x_train = np.load('x_train.npy')
y_train = np.load('y_train.npy')
epochs = 60
batch_size = 64
model.fit(x_train, y_train,
batch_size=batch_size, epochs=epochs,
verbose=1, callbacks=callbacks,
validation_split=0.2)
pred1 = model.predict(x_train, batch_size=1, verbose=1)
# 计算IoU
xx1 = np.maximum(pred1[:, 0], Y_train[:, 0])
yy1 = np.maximum(pred1[:, 1], Y_train[:, 1])
xx2 = np.minimum(pred1[:, 2], Y_train[:, 2])
yy2 = np.minimum(pred1[:, 3], Y_train[:, 3])
area1 = (pred1[:, 2] - pred1[:, 0]) * (pred1[:, 3] - pred1[:, 1])
area2 = (Y_train[:, 2] - Y_train[:, 0]) * (Y_train[:, 3] - Y_train[:, 1])
inter_area = (np.maximum(0, xx2 - xx1)) * (np.maximum(0, yy2 - yy1))
iou = inter_area / (area1 + area2 - inter_area + 1e-6)
k = 0
for j in range(2000):
if iou[j] >= 0.5:
k += 1
acc = k / 2000
print(acc)
#!/usr/bin/env python
from scipy.misc import imread
from xml.etree import ElementTree as ET
# 保存结果到xml文件
for i, img_name in enumerate(img_names):
img_read = imread(img_path + img_name)
# Parse the outputs.
det_xmin = results[i][0]
det_ymin = results[i][1]
det_xmax = results[i][2]
det_ymax = results[i][3]
# 指定根节点/子节点
prediction = ET.Element("prediction")
bndbox = ET.SubElement(prediction, "bndbox")
xmin = ET.SubElement(bndbox, "xmin")
ymin = ET.SubElement(bndbox, "ymin")
xmax = ET.SubElement(bndbox, "xmax")
ymax = ET.SubElement(bndbox, "ymax")
conf = ET.SubElement(bndbox, "conf")
if det_xmax > img_read.shape[1]:
det_xmax = img_read.shape[1] - 1
if det_ymax > img_read.shape[0]:
det_ymax = img_read.shape[0] - 1
xmin.text = str(int(round(det_xmin)))
ymin.text = str(int(round(det_ymin)))
xmax.text = str(int(round(det_xmax)))
ymax.text = str(int(round(det_ymax)))
conf.text = '1.0000'
# 保存结构树
tree = ET.ElementTree(prediction)
tree.write(box_path + img_name[:-3] + 'xml')
#!/usr/bin/env python
from PIL import Image
import numpy as np
import xml.etree.ElementTree as ET
# 获取文件名
import os
all_file = []
def get_name(path):
for file_name in os.listdir(path):
all_file.append(file_name[:-4])
path = 'trainval/img/'
get_name(path)
print(all_file)
num_image_train = len(all_file)
# 导入训练数据
x_train = np.zeros((num_image_train, 224, 224, 3))
box = []
img_size = []
img_path = 'trainval/img/'
xml_path = 'trainval/box/'
for i in range(num_image_train):
img = Image.open(img_path + all_file[i] + '.jpg')
img_size.append(img.size)
img_resize = img.resize((224, 224))
x_train[i] = np.array(img_resize)
xml = ET.parse(xml_path + all_file[i] + '.xml')
box.append([int(xml.find(tag).text)-1
for tag in ['xmin', 'ymin', 'xmax', 'ymax']])
#box[i] = [x / img_size[i][index % 2] for index, x in enumerate(box[i])]
print("训练数据导入完毕")
# 存储数据
np.save("x_train.npy", x_train)
np.save("y_train.npy", np.array(box))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment