Tensorflow 2 图像批量预测返回结果

Posted

技术标签:

【中文标题】Tensorflow 2 图像批量预测返回结果【英文标题】:Tensorflow 2 Image Batch Prediction Return Results 【发布时间】:2021-01-15 11:04:46 【问题描述】:

我有一个已经训练好的模型,我想对目录中的图像进行二元分类预测。我有超过 100,000 张图像,所以为了提高效率,我想做批量预测。如何对我的图像进行批量预测,获取预测结果,并将图像根据类预测后存储在两个单独的文件夹中?

这是我的代码到目前为止的样子...

model_filepath = r"C:\Users\model_200.h5"
model = tf.keras.models.load_model(model_filepath)

test_dir = r"C:\Users\image_testing_folder"

batch_size = 64
IMG_HEIGHT = 200
IMG_WIDTH = 200

test_image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255)
test_image_gen = test_image_generator.flow_from_directory(directory=str(test_dir),
                                                         batch_size=batch_size,
                                                         shuffle=False,
                                                         target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                         )

predictions = (model.predict(test_image_gen) > 0.5).astype("int32")
predictions

一种解决方案是将预测与图像文件路径联系起来,然后使用 shutil.move() 将原始图像移动到目标文件夹。我该怎么做?有没有比使用 ImageDataGenerator 和 .flow_from_directory 更好的批量预测方法?

【问题讨论】:

【参考方案1】:

你可以做一个自定义数据集,这样你也可以轻松地检索文件名:

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras import Sequential
from glob2 import glob
from shutil import copy
import numpy as np

files = glob('group1\\*\\*.jpg')

imsize = 64

def load(file_path):
    img = tf.io.read_file(file_path)
    img = tf.image.decode_png(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)
    img = tf.image.resize(img, size=(imsize, imsize))
    return img, file_path

ds = tf.data.Dataset.from_tensor_slices(files).\
    take(100).\
    shuffle(100).\
    map(load).batch(4)

model = Sequential()
model.add(Conv2D(8, (3, 3), input_shape=(imsize, imsize, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(units=32, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(units=2, activation='sigmoid'))
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

model.build(input_shape=(imsize, imsize, 3))

categories = np.array(['cats', 'dogs'])

target_dir = 'newpics'

for cat in categories:
    os.makedirs(os.path.join(target_dir, cat), exist_ok=True)

for images, filenames in ds:
    preds = model(images)
    targets = categories[np.argmax(preds, axis=1)]
    for file, destination in zip(filenames, targets):
        copy(file.numpy().decode(), os.path.join(target_dir, destination,
                                os.path.basename(file.numpy().decode())
                                ))
        print(file.numpy().decode(), '-->', os.path.join(target_dir, destination,
                                os.path.basename(file.numpy().decode())
                                ))
group1\cats\cat.4051.jpg --> newpics\cats\cat.4051.jpg
group1\cats\cat.4091.jpg --> newpics\dogs\cat.4091.jpg
group1\cats\cat.4055.jpg --> newpics\cats\cat.4055.jpg
group1\cats\cat.4041.jpg --> newpics\cats\cat.4041.jpg
group1\cats\cat.4090.jpg --> newpics\cats\cat.4090.jpg
group1\cats\cat.4071.jpg --> newpics\dogs\cat.4071.jpg
group1\cats\cat.4082.jpg --> newpics\cats\cat.4082.jpg
group1\cats\cat.4037.jpg --> newpics\cats\cat.4037.jpg
group1\cats\cat.4005.jpg --> newpics\cats\cat.4005.jpg

您只需更改 glob 模式和文件夹。

【讨论】:

感谢您的快速回复! 一个问题;尽管您创建了一个自定义数据集,每个“批次”包含 100 张图像,但您使用 for 循环遍历每个图像。从“一次”为一组图像生成预测的意义上说,您的解决方案是真正的“批量预测”吗?它在预测速度方面如何与 model.predict_classes(np.vstack(test_images_2)) 或使用 TF ImageDataGenerator 之类的解决方案进行比较?也许 for 循环是返回文件路径所必需的。谢谢! 对于第一个,您需要将所有图像加载到内存中,这很可能会导致内存过载。对于第二个,它是相同的,因为 ImageDataGenerator 在后台使用了我的实现。正如您在上面看到的,它使用 4 的批量大小,因此一次可以预测多个样本。所以确实更快。你甚至可以提高这个数字。如果它适合内存,您可以使用批量大小 256。

以上是关于Tensorflow 2 图像批量预测返回结果的主要内容,如果未能解决你的问题,请参考以下文章

MobileNet实战:tensorflow2.X版本,MobileNetV3图像分类任务(小数据集)

Tensorflow 模型返回恒定预测 [关闭]

神经分割网络根据测试批量大小给出不同的输出

根据模型预测,使用丢失的自定义函数的返回值

无法预测 ImageNet 类的输入 [Keras + Tensorflow]

基于tensorflow的简单鼠标键盘识别