keras语义分割--从保存的模型中分割测试图像

训练模型

参照keras官方例程得到训练模型，将其保存在路径：save_model\oxford_segmentation.h5，
即将代码callbacks = [ keras.callbacks.ModelCheckpoint("oxford_segmentation.h5", save_best_only=True)]
改为callbacks = [ keras.callbacks.ModelCheckpoint("save_model\oxford_segmentation.h5", save_best_only=True)]，改变仅仅为将模型保存在save_model文件夹下。

从路径中读取一张图像并分割(不保存结果，仅显示)

import cv2 as cv
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt


#  加载训练好的模型
model = keras.models.load_model(filepath=r"save_model\oxford_segmentation.h5")

#  从路径中读取测试图像
img = cv.imread(r"The Oxford-IIIT Pet Dataset\2test\imgs\1\Abyssinian_224.jpg")
#  使图像满足模型预测的要求，即[batch_size, 模型要求的高，模型要求的宽，模型要求的通道数]
img = cv.resize(img[:, :, ::-1], (160, 160))
batch_test_img = np.expand_dims(img, axis=0)
#  预测
pred = model.predict(batch_test_img)

#  取出预测结果，由于只有一张，因此索引为0
pred = pred[0]
#  由于预测结果是像素概率，因此需要将概率最大的类别取出，得到索引图并转换为uint8格式
pred = np.argmax(pred, axis=-1)
pred = pred.astype(dtype=np.uint8)

#  最后将索引图转换为灰度图，并显示
rst = pred.copy()
cv.normalize(pred, rst, alpha=0, beta=255, norm_type=cv.NORM_MINMAX)
plt.subplot(121)
plt.imshow(img)
plt.subplot(122)
plt.imshow(rst, cmap='gray')
plt.show()

从生成器中读取图像并分割(保存文件名是序号)

import cv2 as cv
from tensorflow import keras
import numpy as np
import os


#  加载训练好的模型
model = keras.models.load_model(filepath=r"save_model\oxford_segmentation.h5")

#  图像数据生成器，无需增强，将所有图像作为一个批次
datagen = keras.preprocessing.image.ImageDataGenerator()
img_size = (160, 160)
batch_size = len(os.listdir(r"The Oxford-IIIT Pet Dataset\2test\imgs\1"))
seed = 1
testgen_imgs = datagen.flow_from_directory(
                directory=r"The Oxford-IIIT Pet Dataset\2test\imgs",
                target_size=img_size,
                class_mode=None,
                batch_size=batch_size,
                seed=seed
)
testgen_labels = datagen.flow_from_directory(
                directory=r"The Oxford-IIIT Pet Dataset\2test\labels",
                target_size=img_size,
                class_mode=None,
                batch_size=batch_size,
                seed=seed
)
#  迭代测试图像生成器，进行预测
for batch_test_imgs, batch_test_labels in zip(testgen_imgs, testgen_labels):
    preds = model.predict(batch_test_imgs)


    def mask2gray(mask, input_type=None):
        """
        :param mask: gt的shape是[h,w], pred的shape是[h,w,c]
        :param input_type: pred需要多处理一步
        :return: 采用最大最小值归一化返回灰度图，单通道
        """
        if input_type is 'pred':
            mask = np.argmax(mask, axis=-1)
        mask = mask.astype(dtype=np.uint8)
        rst = mask.copy()
        cv.normalize(mask, rst, alpha=0, beta=255, norm_type=cv.NORM_MINMAX)
        return rst

    for i in range(batch_test_imgs.shape[0]):
        img = batch_test_imgs[i]
        img = img.astype(dtype=np.uint8)

        gt = batch_test_labels[i, :, :, 1]
        gt = mask2gray(gt)

        pred = preds[i]
        pred = mask2gray(pred, input_type='pred')

        save_img = np.zeros([img_size[0], img_size[1]*3, 3], dtype=np.uint8)
        save_img[:, 0:img_size[0], :] = img[:, :, ::-1]
        save_img[:, img_size[0]:img_size[1]*2, :] = cv.cvtColor(gt, cv.COLOR_GRAY2RGB)
        save_img[:, img_size[0]*2:img_size[1]*3, :] = cv.cvtColor(pred, cv.COLOR_GRAY2RGB)
        cv.imwrite(r"The Oxford-IIIT Pet Dataset\2test\preds\{0}.png".format(str(i)), save_img)
    #  由于这个生成器是无限循环迭代，因此，在这个批次(共一个批次)迭代完成后就停止
    break
print("Successfully Save")

从路径中读取图像并分割(保存文件名是原图名)

import cv2 as cv
from tensorflow import keras
import numpy as np
import os


def load_image_from_directory(images_dir, img_size, suffix='.jpg', dtype=np.float32):
    """
    从数据集目录中加载图像数组
    :param images_dir: 数据集目录，原图或标签，该文件夹下直接是图像，三通道，24位深度
    :param img_size: 网络要求的图像大小
    :param suffix: 图像后缀
    :param dtype: 图像数据类型
    :return: 返回图像数组，[图像个数，高，宽，通道数]
    """
    images_path = []
    for fname in os.listdir(images_dir):
        if fname.endswith(suffix) and not fname.startswith('.'):
            images_path.append(os.path.join(images_dir, fname))
    images_path = sorted(images_path)
    images_num = len(os.listdir(images_dir))
    images = np.zeros((images_num,)+img_size+(3,), dtype=dtype)
    for i, path in enumerate(images_path):
        img = cv.imread(path)
        if img.shape[:2] is not img_size:
            img = cv.resize(img, dsize=img_size)
        img = img[:, :, ::-1]
        images[i] = img
    return images, images_path


#  加载训练好的模型
model = keras.models.load_model(filepath=r"save_model\oxford_segmentation.h5")

img_size = (160, 160)
test_images, test_images_path = load_image_from_directory(images_dir=r"The Oxford-IIIT Pet Dataset\2test\imgs\1", img_size=img_size)
test_labels, _ = load_image_from_directory(images_dir=r"The Oxford-IIIT Pet Dataset\2test\labels\1", img_size=img_size, suffix='.png')
preds = model.predict(test_images)


def mask2gray(mask, input_type=None):
    """
    :param mask: gt的shape是[h,w], pred的shape是[h,w,c]
    :param input_type: pred需要多处理一步
    :return: 采用最大最小值归一化返回灰度图，单通道
    """
    if input_type is 'pred':
        mask = np.argmax(mask, axis=-1)
    mask = mask.astype(dtype=np.uint8)
    rst = mask.copy()
    cv.normalize(mask, rst, alpha=0, beta=255, norm_type=cv.NORM_MINMAX)
    return rst


for i in range(test_images.shape[0]):
    img = test_images[i]
    img = img.astype(dtype=np.uint8)

    gt = test_labels[i, :, :, 1]
    gt = mask2gray(gt)

    pred = preds[i]
    pred = mask2gray(pred, input_type='pred')

    save_img = np.zeros([img_size[0], img_size[1]*3, 3], dtype=np.uint8)
    save_img[:, 0:img_size[0], :] = img[:, :, ::-1]
    save_img[:, img_size[0]:img_size[1]*2, :] = cv.cvtColor(gt, cv.COLOR_GRAY2RGB)
    save_img[:, img_size[0]*2:img_size[1]*3, :] = cv.cvtColor(pred, cv.COLOR_GRAY2RGB)
    cv.imwrite(r"The Oxford-IIIT Pet Dataset\2test\preds\{0}".format(test_images_path[i].split('\\')[-1]), save_img)