使用 keras 和预训练的 pspnet50 进行图像分割

Question

我有一个预训练的 keras 模型 pspnet50_ad20k 并且想从中得到一个分割图像。输入是形状为 (1, 473, 473, 3) 的 numpy 数组中的图像，它 returns 是形状为 (1, 473, 473, 150) 的数组，因为有 150 个不同类这个模型预测。

import os
os.environ['KERAS_BACKEND'] = "tensorflow"
from keras.models import model_from_json

json_file = open('/data/pspnet50_ade20k.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
loaded_model = model_from_json(loaded_model_json)
# load weights into new model
loaded_model.load_weights("/data/pspnet50_ade20k.h5")
print("Loaded model from disk")

# img_array contains img in shape (1, 473, 473, 3)
result = loaded_model.predict(img_array)
# result contains img in shape (1, 473, 473, 150)

我的问题：如何从生成的数组到分割图像？我不知何故必须为图像中预测的 150 类着色，但我没有不知道怎么办。有人可以给我解释一下吗？

Answer 1

json_file = open('/data/pspnet50_ade20k.json', 'r') 您使用的模型以预测 50 类的方式编码，重新编码更改模型或构建您自己的模型来预测您需要多少类。

最好使用 keras 构建您自己的模型，并加载以下权重：loaded_model.load_weights("/data/pspnet50_ade20k.h5")

Answer 2

Ade20k 数据集有 150 个类，这就是您的网络最终输出 (1, 473, 473, 150) 的原因。 150表示每个像素位置150类的概率。所以需要改最后一层，一般是conv2d(256,classes=150,kernel_size=1)。请将 150 替换为 50，这是您要预测的类个数，并在您自己的数据集上微调您的模型。

对于着色预测，你可以这样做，

from PIL import Image
import numpy as np
prediction = np.argmax(prediction, axis=3) # get the max prob label
palette = np.loadtxt('path/to/ade20k_colors.txt').astype('uint8')
# The format of `ade20k_colors.txt`
# 120 120 120
# 180 120 120
# 6 230 230
# 80 50 50
# 4 200 3
# ...
color = Image.fromarray(prediction.astype(np.uint8)).convert('P')
color.putpalette(palette)
# color is colorized prediction.

使用 keras 和预训练的 pspnet50 进行图像分割

Image segmentation with keras and pretrained pspnet50

numpy

image-segmentation

deep-learning

keras

tensorflow