Keras:Unet网络实现多类语义分割方式

1 介绍

u-net最初是用来对医学图像的语义分割，后来也有人将其应用于其他领域。但大多还是用来进行二分类，即将原始图像分成两个灰度级或者色度，依次找到图像中感兴趣的目标部分。

本文主要利用u-net网络结构实现了多类的语义分割，并展示了部分测试效果，希望对你有用！

2 源代码

（1）训练模型

from __future__ import print_function
import os
import datetime
import numpy as np
from keras.models import model
from keras.layers import input, concatenate, conv2d, maxpooling2d, conv2dtranspose, averagepooling2d, dropout, \
 batchnormalization
from keras.optimizers import adam
from keras.layers.convolutional import upsampling2d, conv2d
from keras.callbacks import modelcheckpoint
from keras import backend as k
from keras.layers.advanced_activations import leakyrelu, relu
import cv2
 
pixel = 512 #set your image size
batch_size = 5
lr = 0.001
epoch = 100
x_channel = 3 # training images channel
y_channel = 1 # label iamges channel
x_num = 422 # your traning data number
 
pathx = 'i:\\pascal voc dataset\\train1\\images\\' #change your file path
pathy = 'i:\\pascal voc dataset\\train1\\segmentationobject\\' #change your file path
 
#data processing
def generator(pathx, pathy,batch_size):
 while 1:
  x_train_files = os.listdir(pathx)
  y_train_files = os.listdir(pathy)
  a = (np.arange(1, x_num))
  x = []
  y = []
  for i in range(batch_size):
   index = np.random.choice(a)
   # print(index)
   img = cv2.imread(pathx + x_train_files[index], 1)
   img = np.array(img).reshape(pixel, pixel, x_channel)
   x.append(img)
   img1 = cv2.imread(pathy + y_train_files[index], 1)
   img1 = np.array(img1).reshape(pixel, pixel, y_channel)
   y.append(img1)
 
  x = np.array(x)
  y = np.array(y)
  yield x, y
 
 #creat unet network
inputs = input((pixel, pixel, 3))
conv1 = conv2d(8, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
pool1 = averagepooling2d(pool_size=(2, 2))(conv1) # 16
 
conv2 = batchnormalization(momentum=0.99)(pool1)
conv2 = conv2d(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
conv2 = batchnormalization(momentum=0.99)(conv2)
conv2 = conv2d(64, 1, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
conv2 = dropout(0.02)(conv2)
pool2 = averagepooling2d(pool_size=(2, 2))(conv2) # 8
 
conv3 = batchnormalization(momentum=0.99)(pool2)
conv3 = conv2d(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
conv3 = batchnormalization(momentum=0.99)(conv3)
conv3 = conv2d(128, 1, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
conv3 = dropout(0.02)(conv3)
pool3 = averagepooling2d(pool_size=(2, 2))(conv3) # 4
 
conv4 = batchnormalization(momentum=0.99)(pool3)
conv4 = conv2d(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
conv4 = batchnormalization(momentum=0.99)(conv4)
conv4 = conv2d(256, 1, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
conv4 = dropout(0.02)(conv4)
pool4 = averagepooling2d(pool_size=(2, 2))(conv4)
 
conv5 = batchnormalization(momentum=0.99)(pool4)
conv5 = conv2d(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)
conv5 = batchnormalization(momentum=0.99)(conv5)
conv5 = conv2d(512, 1, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)
conv5 = dropout(0.02)(conv5)
pool4 = averagepooling2d(pool_size=(2, 2))(conv4)
# conv5 = conv2d(35, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
# drop4 = dropout(0.02)(conv5)
pool4 = averagepooling2d(pool_size=(2, 2))(pool3) # 2
pool5 = averagepooling2d(pool_size=(2, 2))(pool4) # 1
 
conv6 = batchnormalization(momentum=0.99)(pool5)
conv6 = conv2d(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)
 
conv7 = conv2d(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)
up7 = (upsampling2d(size=(2, 2))(conv7)) # 2
conv7 = conv2d(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up7)
merge7 = concatenate([pool4, conv7], axis=3)
 
conv8 = conv2d(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge7)
up8 = (upsampling2d(size=(2, 2))(conv8)) # 4
conv8 = conv2d(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up8)
merge8 = concatenate([pool3, conv8], axis=3)
 
conv9 = conv2d(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge8)
up9 = (upsampling2d(size=(2, 2))(conv9)) # 8
conv9 = conv2d(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up9)
merge9 = concatenate([pool2, conv9], axis=3)
 
conv10 = conv2d(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9)
up10 = (upsampling2d(size=(2, 2))(conv10)) # 16
conv10 = conv2d(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up10)
 
conv11 = conv2d(16, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv10)
up11 = (upsampling2d(size=(2, 2))(conv11)) # 32
conv11 = conv2d(8, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up11)
 
# conv12 = conv2d(3, 1, activation='relu', padding='same', kernel_initializer='he_normal')(conv11)
conv12 = conv2d(3, 1, activation='relu', padding='same', kernel_initializer='he_normal')(conv11)
 
model = model(input=inputs, output=conv12)
print(model.summary())
model.compile(optimizer=adam(lr=1e-3), loss='mse', metrics=['accuracy'])
 
history = model.fit_generator(generator(pathx, pathy,batch_size),
        steps_per_epoch=600, nb_epoch=epoch)
end_time = datetime.datetime.now().strftime('%y-%m-%d %h:%m:%s')
 
 #save your training model
model.save(r'v1_828.h5')
 
#save your loss data
mse = np.array((history.history['loss']))
np.save(r'v1_828.npy', mse)

（2）测试模型

from keras.models import load_model
import numpy as np
import matplotlib.pyplot as plt
import os
import cv2
 
model = load_model('v1_828.h5')
test_images_path = 'i:\\pascal voc dataset\\test\\test_images\\'
test_gt_path = 'i:\\pascal voc dataset\\test\\segmentationobject\\'
pre_path = 'i:\\pascal voc dataset\\test\\pre\\'
 
x = []
for info in os.listdir(test_images_path):
 a = cv2.imread(test_images_path + info)
 x.append(a)
 # i += 1
x = np.array(x)
print(x.shape)
y = model.predict(x)
 
groudtruth = []
for info in os.listdir(test_gt_path):
 a = cv2.imread(test_gt_path + info)
 groudtruth.append(a)
groudtruth = np.array(groudtruth)
 
i = 0
for info in os.listdir(test_images_path):
 cv2.imwrite(pre_path + info,y[i])
 i += 1
 
a = range(10)
n = np.random.choice(a)
cv2.imwrite('prediction.png',y[n])
cv2.imwrite('groudtruth.png',groudtruth[n])
fig, axs = plt.subplots(1, 3)
# cnt = 1
# for j in range(1):
axs[0].imshow(np.abs(x[n]))
axs[0].axis('off')
axs[1].imshow(np.abs(y[n]))
axs[1].axis('off')
axs[2].imshow(np.abs(groudtruth[n]))
axs[2].axis('off')
 # cnt += 1
fig.savefig("imagestest.png")
plt.close()

3 效果展示

说明：从左到右依次是预测图像，真实图像，标注图像。可以看出，对于部分数据的分割效果还有待改进，主要原因还是数据集相对复杂，模型难于找到其中的规律。

以上这篇keras:unet网络实现多类语义分割方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持。