Сегментация изображения с помощью кераса и предварительно обученного pspnet50

Question

У меня есть модель pspnet50_ad20k с предварительно подготовленными керасами, и я хочу получить из нее сегментированное изображение. Входные данные - это изображение в массиве фигур в форме (1, 473, 473, 3), и оно возвращает массив в форме (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

Набор данных 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. 

Answer 2

json_file = open ('/ data / pspnet50_ade20k.json', 'r') используемая вами модель закодирована таким образом, чтобы предсказать 50 классов, перекодировать изменение модели или построить собственную модель, чтобы предсказать, сколько классов вам нужно.

лучше использовать керасы для построения своей собственной модели с такими загруженными весами :loaded_model.load_weights ("/ data / pspnet50_ade20k.h5")