21个tensorflow项目(三):打造自己的图像识别模型
- 环境介绍
- 数据准备
- 制作TFrecord
- 制作TFrecord代码:
- 读取TFrecord代码:
- 训练模型代码
- 代码
- 代码解析
- 运行结果
- 预测代码
- 代码
- 运行结果
- 自己对数据集的一些疑惑
环境介绍
Python版本:Python 3.8.16
TensorFlow版本:2.6.0
数据准备
书中选用的数据集为卫星图像数据集,在这里也采用同样的数据集
下载地址为:https://wwn.lanzout.com/i6KkX0t8ulfe 解压后的如图所示:
类别名 | 含义 | 示例图像 |
glacier | 冰川 | |
rock | 岩石 | |
urban | 城市区域 | |
water | 水域 | |
wetland | 农田 | |
wood | 森林 |
制作TFrecord
在TensorFlow 2.x中,制作TFRecord主要涉及以下步骤:
- 定义特征信息:首先需要定义每个样本的特征信息,包括特征的名称、类型和形状等。可以使用tf.train.Feature类来定义每个特征。
- 将特征信息转换为Example格式:将每个样本的特征信息转换为tf.train.Example格式,可以使用tf.train.Example类来完成。
- 将Example序列化为字符串:将每个tf.train.Example序列化为字符串,可以使用tf.io.serialize_tensor()函数。
- 将序列化后的字符串写入TFRecord文件:使用tf.io.TFRecordWriter类将序列化后的字符串写入TFRecord文件中。
制作TFrecord代码:
import tensorflow as tf
import os
image_labels = {"glacier": 0, "rock": 1, "urban": 2, "water": 3 , "wetland": 4, "wood": 5 }
def _bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def _int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def image_example(image_string, label):
image_shape = tf.image.decode_jpeg(image_string).shape
feature = {
'height': _int64_feature(image_shape[0]),
'width': _int64_feature(image_shape[1]),
'depth': _int64_feature(image_shape[2]),
'label': _int64_feature(label),
'image_raw': _bytes_feature(image_string),
}
return tf.train.Example(features=tf.train.Features(feature=feature))
def creat_tfrecord(path,path_data):
with tf.io.TFRecordWriter(path_data) as writer:
for label in os.listdir(path):
if label not in image_labels:
continue
label_path = os.path.join(path, label)
for image_name in os.listdir(label_path):
image_path = os.path.join(label_path, image_name)
tf_example = image_example(open(image_path, 'rb').read(), image_labels[label])
writer.write(tf_example.SerializeToString())
train_folder = "data_prepare/pic/train"
validation_folder = "data_prepare/pic/validation"
train_output_file = "./train.tfrecord"
validation_output_file = "./validation.tfrecord"
creat_tfrecord(train_folder,train_output_file)
creat_tfrecord(validation_folder,validation_output_file)
读取TFrecord代码:
import tensorflow as tf
raw_image_dataset = tf.data.TFRecordDataset('train.tfrecord')
image_feature_description = {
'height': tf.io.FixedLenFeature([], tf.int64),
'width': tf.io.FixedLenFeature([], tf.int64),
'depth': tf.io.FixedLenFeature([], tf.int64),
'label': tf.io.FixedLenFeature([], tf.int64),
'image_raw': tf.io.FixedLenFeature([], tf.string),
}
def _parse_image_function(example_proto):
return tf.io.parse_single_example(example_proto, image_feature_description)
parsed_image_dataset = raw_image_dataset.map(_parse_image_function)
for raw_record in parsed_image_dataset.take(1):
print(repr(raw_record))
训练模型代码
代码
import tensorflow as tf
from tensorflow.keras import layers, models, optimizers
from tensorflow.keras.applications.inception_v3 import InceptionV3
from tensorflow.keras.preprocessing.image import ImageDataGenerator
NUM_CLASSES = 6
IMG_SIZE = (299, 299)
BATCH_SIZE = 32
LEARNING_RATE = 0.001
EPOCHS = 20
TRAIN_TFRECORD = "./train.tfrecord"
VALID_TFRECORD = "./validation.tfrecord"
def parse_tfrecord(serialized_example):
feature_description = {
'height': tf.io.FixedLenFeature([], tf.int64),
'width': tf.io.FixedLenFeature([], tf.int64),
'depth': tf.io.FixedLenFeature([], tf.int64),
'label': tf.io.FixedLenFeature([], tf.int64),
'image_raw': tf.io.FixedLenFeature([], tf.string),
}
example = tf.io.parse_single_example(serialized_example, feature_description)
image = tf.image.decode_jpeg(example['image_raw'], channels=3)
image = tf.image.resize(image, IMG_SIZE)
image = tf.cast(image, tf.float32) / 255.0
label = tf.one_hot(example['label'], NUM_CLASSES)
return image, label
def create_dataset(file_path):
dataset = tf.data.TFRecordDataset(file_path)
dataset = dataset.map(parse_tfrecord, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.shuffle(1000)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(tf.data.AUTOTUNE)
return dataset
train_dataset = create_dataset(TRAIN_TFRECORD)
valid_dataset = create_dataset(VALID_TFRECORD)
base_model = InceptionV3(weights='imagenet', include_top=False, input_shape=(*IMG_SIZE, 3))
x = layers.GlobalAveragePooling2D()(base_model.output)
x = layers.Dense(1024, activation='relu')(x)
x = layers.Dropout(0.5)(x)
x = layers.Dense(NUM_CLASSES, activation='softmax')(x)
model = models.Model(inputs=base_model.input, outputs=x)
for layer in base_model.layers:
layer.trainable = False
model.compile(optimizer=optimizers.Adam(LEARNING_RATE),
loss='categorical_crossentropy',
metrics=['accuracy'])
history = model.fit(train_dataset, epochs=EPOCHS, validation_data=valid_dataset)
loss, accuracy = model.evaluate(valid_dataset)
print("Validation accuracy:", accuracy)
model.save("13.h5")
代码解析
- 定义常量和超参数:NUM_CLASSES 表示图像分类的类别数,IMG_SIZE 表示输入图像的尺寸,BATCH_SIZE 表示每个 batch 的大小,LEARNING_RATE 表示优化器的学习率,EPOCHS 表示模型训练的轮数,TRAIN_TFRECORD 和 VALID_TFRECORD 表示训练和验证数据集的 TFRecord 文件路径。
- 定义用于读取和预处理 TFRecord 数据的函数 parse_tfrecord,其中通过 tf.io.parse_single_example 解析 TFRecord 文件,然后将图像进行解码、缩放和类型转换,将标签进行 one-hot 编码,并返回图像和标签。
- 定义用于创建数据集的函数 create_dataset,其中通过 tf.data.TFRecordDataset 读取 TFRecord 文件,然后使用 map 函数应用 parse_tfrecord 函数对每个样本进行解析和预处理,再对数据进行 shuffle、batch 和 prefetch 操作,并返回处理后的数据集。
- 创建 train_dataset 和 valid_dataset 分别对应训练和验证数据集。
- 定义 InceptionV3 模型,使用 InceptionV3 函数从预训练的 InceptionV3 模型中加载权重,同时指定 include_top=False 表示只加载卷积部分的权重,将其输出作为新模型的输入。
- 在 InceptionV3 模型的输出上添加新的全局平均池化层、全连接层和 Dropout 层,并输出 NUM_CLASSES 个类别的概率分布。
- 构建模型,指定输入和输出,并将 InceptionV3 模型的层冻结。
- 使用 Adam 优化器、交叉熵损失函数和准确率指标编译模型。
- 对训练数据集进行 EPOCHS 轮训练,使用验证数据集评估模型性能。
- 最后将训练好的模型保存为 13.h5 文件。
运行结果
我使用20个epoch来训练,每个epoch包含150个batch。训练过程中,损失函数的值从1.0670下降到了0.3733,准确率从0.7135上升到了0.8617。同时也可以看到,验证集的损失值从1.0338下降到了0.6949,准确率从0.6058上升到了0.7692。模型的训练过程看起来还不错,但我觉得不是很好。
预测代码
代码
import tensorflow as tf
import numpy as np
from PIL import Image
NUM_CLASSES = 6
IMG_SIZE = (299, 299)
model = tf.keras.models.load_model('13.h5')
def preprocess_image(image_path):
img = Image.open(image_path)
img = img.resize(IMG_SIZE)
img = np.array(img) / 255.0
img = np.expand_dims(img, axis=0)
return img
image_path = 'data_prepare/pic/validation/wood/73987_98782_18.jpg'
preprocessed_image = preprocess_image(image_path)
prediction = model.predict(preprocessed_image)
predicted_class = np.argmax(prediction)
print("Predicted class:", predicted_class)
运行结果
可以看出正确的预测出了图片的类别
自己对数据集的一些疑惑
本次的数据集我是从网上下载下来的数据集与书本上的一致,不知道是我的下载地址有误,还是数据集本身是这样的,我觉得数据集有些图片的标签是不是标错了例如:
在数据集中这张图片的标签为wetland,但这明显根本不是wetland,而是urban。还有一部分的图片也有这样的情况,不知道是什么原因。所以我认为模型训练的效果不理想可能跟这个有关,当然,这只是我自己的看法。