一、代码中的数据集可以点击以下链接进行下载
二、代码运行环境
Tensorflow-gpu==2.4.0
Python==3.7
三、数据集处理的代码如下所示
import os
import random
import tensorflow as tf
import xml.etree.ElementTree as ET
import glob
import numpy as np
import matplotlib.pyplot as plt
# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
def load_bind_box(xml_path):
xml_file = open(xml_path, encoding='utf-8')
tree = ET.parse(xml_file)
root = tree.getroot()
width = float(root.find('size').find('width').text)
height = float(root.find('size').find('height').text)
xmin = float(root.find('object').find('bndbox').find('xmin').text) / width
ymin = float(root.find('object').find('bndbox').find('ymin').text) / height
xmax = float(root.find('object').find('bndbox').find('xmax').text) / width
ymax = float(root.find('object').find('bndbox').find('ymax').text) / height
return [xmin, ymin, xmax, ymax]
def load_image(image_path):
img = tf.io.read_file(image_path)
img = tf.image.decode_jpeg(img, channels=3)
img = tf.image.resize(img, [224, 224])
img = tf.cast(img, tf.float32)
img = img / 255
return img
def make_dataset():
labels_path = glob.glob(r'dataset\annotations\*.xml')
random.shuffle(labels_path)
images_path = [os.path.join(r'dataset\images', i.split('\\')[-1].split('.')[0]) + '.jpg' for i in labels_path]
labels_box = [load_bind_box(i) for i in labels_path]
out1, out2, out3, out4 = list(zip(*labels_box))
out1 = np.array(out1)
out2 = np.array(out2)
out3 = np.array(out3)
out4 = np.array(out4)
label_dataset = tf.data.Dataset.from_tensor_slices((out1, out2, out3, out4))
image_dataset = tf.data.Dataset.from_tensor_slices(images_path)
image_dataset = image_dataset.map(load_image)
count = len(labels_path)
train_count = int(count * 0.8)
test_count = count - train_count
dataset = tf.data.Dataset.zip((image_dataset, label_dataset))
train_dataset = dataset.take(train_count)
test_dataset = dataset.skip(train_count)
train_dataset = train_dataset.repeat().shuffle(train_count).batch(32)
test_dataset = test_dataset.batch(32)
return train_dataset, train_count, test_dataset, test_count
if __name__ == '__main__':
data, num, test, test_count = make_dataset()
for img, label in data.take(1):
img = img[0]
out1, out2, out3, out4 = label
xmin, ymin, xmax, ymax = out1[0].numpy(), out2[0].numpy(), out3[0].numpy(), out4[0].numpy()
my_label = [ymin, xmin, ymax, xmax]
my_label = tf.cast(my_label, tf.float32)
my_label = tf.expand_dims(tf.expand_dims(my_label, 0), 0)
img = tf.expand_dims(img, 0)
img = tf.image.draw_bounding_boxes(img, my_label, [[1.0, 0.0, 0.0]])
plt.imshow(img[0])
plt.show()
四、模型的构建代码如下所示
import tensorflow as tf
import os
# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
# 进行模型的构建
def make_model():
xception = tf.keras.applications.Xception(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
inputs = tf.keras.layers.Input(shape=(224, 224, 3))
x = xception(inputs)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(2048, activation='relu')(x)
x = tf.keras.layers.Dense(256, activation='relu')(x)
out1 = tf.keras.layers.Dense(1)(x)
out2 = tf.keras.layers.Dense(1)(x)
out3 = tf.keras.layers.Dense(1)(x)
out4 = tf.keras.layers.Dense(1)(x)
prediction = [out1, out2, out3, out4]
model = tf.keras.models.Model(inputs=inputs, outputs=prediction)
model.summary()
return model
if __name__ == '__main__':
my_model = make_model()
五、模型的训练代码如下所示
import os
import tensorflow as tf
from data_loader import make_dataset
from model_loader import make_model
# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
# 加载数据集
train_dataset, train_count, test_dataset, test_count = make_dataset()
# 加载模型
model = make_model()
# 模型的相关配置
model.compile(
tf.keras.optimizers.Adam(learning_rate=0.0001),
loss='mse',
metrics=['mae']
)
# 模型的训练
model.fit(train_dataset,
epochs=50,
steps_per_epoch=train_count // 32,
validation_steps=test_count // 32,
validation_data=test_dataset
)
# 模型的保存
model.save(r'model_data/model.h5')
六、模型的预测代码如下所示
import os
import tensorflow as tf
from data_loader import load_image
import matplotlib.pyplot as plt
# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
# 模型的加载
model = tf.keras.models.load_model(r'model_data/model.h5')
# 开始进行预测
while 1:
path = input('请输入图片的路径:')
image = load_image(path)
image = tf.expand_dims(image, axis=0)
xmin, ymin, xmax, ymax = model.predict(image)
bind_box = tf.expand_dims(tf.expand_dims([ymin[0][0], xmin[0][0], ymax[0][0], xmax[0][0]], axis=0), axis=0)
image = tf.image.draw_bounding_boxes(image, bind_box, [[1.0, 0.0, 0.0]])
plt.imshow(image[0])
plt.show()