Keras:Unet網路實現多類語義分割方式
阿新 • • 發佈:2020-06-11
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,Y_CHANNEL) Y.append(img1) X = np.array(X) Y = np.array(Y) yield X,Y #creat unet network inputs = Input((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,kernel_initializer='he_normal')(conv2) conv2 = BatchNormalization(momentum=0.99)(conv2) conv2 = Conv2D(64,1,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,kernel_initializer='he_normal')(conv3) conv3 = BatchNormalization(momentum=0.99)(conv3) conv3 = Conv2D(128,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,kernel_initializer='he_normal')(conv4) conv4 = BatchNormalization(momentum=0.99)(conv4) conv4 = Conv2D(256,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,kernel_initializer='he_normal')(conv5) conv5 = BatchNormalization(momentum=0.99)(conv5) conv5 = Conv2D(512,kernel_initializer='he_normal')(conv5) conv5 = Dropout(0.02)(conv5) pool4 = AveragePooling2D(pool_size=(2,2))(conv4) # conv5 = Conv2D(35,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,kernel_initializer='he_normal')(conv6) conv7 = Conv2D(256,kernel_initializer='he_normal')(conv6) up7 = (UpSampling2D(size=(2,2))(conv7)) # 2 conv7 = Conv2D(256,kernel_initializer='he_normal')(up7) merge7 = concatenate([pool4,conv7],axis=3) conv8 = Conv2D(128,kernel_initializer='he_normal')(merge7) up8 = (UpSampling2D(size=(2,2))(conv8)) # 4 conv8 = Conv2D(128,kernel_initializer='he_normal')(up8) merge8 = concatenate([pool3,conv8],axis=3) conv9 = Conv2D(64,kernel_initializer='he_normal')(merge8) up9 = (UpSampling2D(size=(2,2))(conv9)) # 8 conv9 = Conv2D(64,kernel_initializer='he_normal')(up9) merge9 = concatenate([pool2,conv9],axis=3) conv10 = Conv2D(32,kernel_initializer='he_normal')(merge9) up10 = (UpSampling2D(size=(2,2))(conv10)) # 16 conv10 = Conv2D(32,kernel_initializer='he_normal')(up10) conv11 = Conv2D(16,kernel_initializer='he_normal')(conv10) up11 = (UpSampling2D(size=(2,2))(conv11)) # 32 conv11 = Conv2D(8,kernel_initializer='he_normal')(up11) # conv12 = Conv2D(3,kernel_initializer='he_normal')(conv11) conv12 = Conv2D(3,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,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網路實現多類語義分割方式就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。