Tensorflow---使用Tensorflow进行图像目标定位的实现

一、代码中的数据集可以点击以下链接进行下载

百度网盘提取码：lala

二、代码运行环境

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()

七、代码的运行结果如下所示