U-net網絡實現醫學圖像分割以及遙感圖像分割源代碼

U-net網絡主要思路是源於FCN，采用全卷積網絡，對圖像進行逐像素分類，能在圖像分割領域達到不錯的效果。

因其網絡結構類似於U型，所以以此命名，可以由其架構清晰的看出，其構成是由左端的卷積壓縮層，以及右端的轉置卷積放大層組成；

左右兩端之間還有聯系，通過灰色箭頭所指，右端在進行轉置卷積操作的時候，會拼接左端前幾次卷積后的結果，這樣可以保證得到

更多的信息。

在網絡的末端得到兩張feature map之后還需要通過softmax函數得到概率圖，整個網絡輸出的是類別數量的特征圖，最后得到的是類別的概率圖，針對每個像素點給出其屬於哪個類別的概率。

Unet網絡的結構提出主要是為了使用在醫學圖像分割領域，它與fcn相比，有以下有優點：1.多尺度  2.能適應超大圖像的

輸入。

unet的卷積過程，是從高分辨率（淺層特征）到低分辨率（深層特征）的過程。

unet的特點就是通過反卷積過程中的拼接，使得淺層特征和深層特征結合起來。對於醫學圖像來說，unet能用深層特征用於定位，淺層特征用於精確分割，所以unet常見於很多圖像分割任務。

 

 

其在Keras實現的部分代碼解析如下：

from model import *
from data import *#導入這兩個文件中的所有函數
 
#os.environ[“CUDA_VISIBLE_DEVICES”] = “0”
 
 
data_gen_args = dict(rotation_range=0.2,
                    width_shift_range=0.05,
                    height_shift_range=0.05,
                    shear_range=0.05,
                    zoom_range=0.05,
                    horizontal_flip=True,
                    fill_mode=‘nearest’)#數據增強時的變換方式的字典
myGene = trainGenerator(2,‘data/membrane/train’,‘image’,‘label’,data_gen_args,save_to_dir = None)
#得到一個生成器，以batch=2的速率無限生成增強后的數據
 
model = unet()
model_checkpoint = ModelCheckpoint(‘unet_membrane.hdf5’, monitor=‘loss’,verbose=1, save_best_only=True)
#回調函數，第一個是保存模型路徑，第二個是檢測的值，檢測Loss是使它最小，第三個是只保存在驗證集上性能最好的模型
 
model.fit_generator(myGene,steps_per_epoch=300,epochs=1,callbacks=[model_checkpoint])
#steps_per_epoch指的是每個epoch有多少個batch_size，也就是訓練集總樣本數除以batch_size的值
#上面一行是利用生成器進行batch_size數量的訓練，樣本和標簽通過myGene傳入
testGene = testGenerator(“data/membrane/test”)
results = model.predict_generator(testGene,30,verbose=1)
#30是step,steps: 在停止之前，來自 generator 的總步數 (樣本批次)。 可選參數 Sequence：如果未指定，將使用len(generator) 作為步數。
#上面的返回值是：預測值的 Numpy 數組。
saveResult(“data/membrane/test”,results)#保存結果
1
data.py文件：

from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
import numpy as np 
import os
import glob
import skimage.io as io
import skimage.transform as trans
 
Sky = [128,128,128]
Building = [128,0,0]
Pole = [192,192,128]
Road = [128,64,128]
Pavement = [60,40,222]
Tree = [128,128,0]
SignSymbol = [192,128,128]
Fence = [64,64,128]
Car = [64,0,128]
Pedestrian = [64,64,0]
Bicyclist = [0,128,192]
Unlabelled = [0,0,0]
 
COLOR_DICT = np.array([Sky, Building, Pole, Road, Pavement,
                          Tree, SignSymbol, Fence, Car, Pedestrian, Bicyclist, Unlabelled])
 
 
def adjustData(img,mask,flag_multi_class,num_class):
    if(flag_multi_class):#此程序中不是多類情況，所以不考慮這個
        img = img / 255
        mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]
#if else的簡潔寫法，一行表達式，為真時放在前面，不明白mask.shape=4的情況是什么，由於有batch_size，所以mask就有3維[batch_size,wigth,heigh],估計mask[:,:,0]是寫錯了，應該寫成[0,:,:],這樣可以得到一片圖片，
        new_mask = np.zeros(mask.shape + (num_class,))
#np.zeros里面是shape元組，此目的是將數據厚度擴展到num_class層，以在層的方向實現one-hot結構
 
        for i in range(num_class):
            #for one pixel in the image, find the class in mask and convert it into one-hot vector
            #index = np.where(mask == i)
            #index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)
            #new_mask[index_mask] = 1
            new_mask[mask == i,i] = 1#將平面的mask的每類，都單獨變成一層，
        new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))
        mask = new_mask
    elif(np.max(img) > 1):
        img = img / 255
        mask = mask /255
        mask[mask > 0.5] = 1
        mask[mask <= 0.5] = 0
    return (img,mask)
#上面這個函數主要是對訓練集的數據和標簽的像素值進行歸一化
 
 
def trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = "grayscale",
                    mask_color_mode = "grayscale",image_save_prefix  = "image",mask_save_prefix  = "mask",
                    flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):
    '''
    can generate image and mask at the same time
    use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same
    if you want to visualize the results of generator, set save_to_dir = "your path"
    '''
    image_datagen = ImageDataGenerator(**aug_dict)
    mask_datagen = ImageDataGenerator(**aug_dict)
    image_generator = image_datagen.flow_from_directory(#https://blog.csdn.net/nima1994/article/details/80626239
        train_path,#訓練數據文件夾路徑
        classes = [image_folder],#類別文件夾,對哪一個類進行增強
        class_mode = None,#不返回標簽
        color_mode = image_color_mode,#灰度，單通道模式
        target_size = target_size,#轉換后的目標圖片大小
        batch_size = batch_size,#每次產生的（進行轉換的）圖片張數
        save_to_dir = save_to_dir,#保存的圖片路徑
        save_prefix  = image_save_prefix,#生成圖片的前綴，僅當提供save_to_dir時有效
        seed = seed)
    mask_generator = mask_datagen.flow_from_directory(
        train_path,
        classes = [mask_folder],
        class_mode = None,
        color_mode = mask_color_mode,
        target_size = target_size,
        batch_size = batch_size,
        save_to_dir = save_to_dir,
        save_prefix  = mask_save_prefix,
        seed = seed)
    train_generator = zip(image_generator, mask_generator)#組合成一個生成器
    for (img,mask) in train_generator:
#由於batch是2，所以一次返回兩張，即img是一個2張灰度圖片的數組，[2,256,256]
        img,mask = adjustData(img,mask,flag_multi_class,num_class)#返回的img依舊是[2,256,256]
        yield (img,mask)
#每次分別產出兩張圖片和標簽，不懂yield的請看https://blog.csdn.net/mieleizhi0522/article/details/82142856
 # yield相當於生成器
#上面這個函數主要是產生一個數據增強的圖片生成器，方便后面使用這個生成器不斷生成圖片
 
 
def testGenerator(test_path,num_image = 30,target_size = (256,256),flag_multi_class = False,as_gray = True):
    for i in range(num_image):
        img = io.imread(os.path.join(test_path,"%d.png"%i),as_gray = as_gray)
        img = img / 255
        img = trans.resize(img,target_size)
        img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img
        img = np.reshape(img,(1,)+img.shape)
#將測試圖片擴展一個維度，與訓練時的輸入[2,256,256]保持一致
        yield img
 
#上面這個函數主要是對測試圖片進行規范，使其尺寸和維度上和訓練圖片保持一致
 
def geneTrainNpy(image_path,mask_path,flag_multi_class = False,num_class = 2,image_prefix = "image",mask_prefix = "mask",image_as_gray = True,mask_as_gray = True):
    image_name_arr = glob.glob(os.path.join(image_path,"%s*.png"%image_prefix))
#相當於文件搜索，搜索某路徑下與字符匹配的文件https://blog.csdn.net/u010472607/article/details/76857493/
    image_arr = []
    mask_arr = []
    for index,item in enumerate(image_name_arr):#enumerate是枚舉，輸出[(0,item0),(1,item1),(2,item2)]
        img = io.imread(item,as_gray = image_as_gray)
        img = np.reshape(img,img.shape + (1,)) if image_as_gray else img
        mask = io.imread(item.replace(image_path,mask_path).replace(image_prefix,mask_prefix),as_gray = mask_as_gray)
#重新在mask_path文件夾下搜索帶有mask字符的圖片（標簽圖片）
        mask = np.reshape(mask,mask.shape + (1,)) if mask_as_gray else mask
        img,mask = adjustData(img,mask,flag_multi_class,num_class)
        image_arr.append(img)
        mask_arr.append(mask)
    image_arr = np.array(image_arr)
    mask_arr = np.array(mask_arr)#轉換成array
    return image_arr,mask_arr
#該函數主要是分別在訓練集文件夾下和標簽文件夾下搜索圖片，然后擴展一個維度后以array的形式返回，是為了在沒用數據增強時的讀取文件夾內自帶的數據
 
 
def labelVisualize(num_class,color_dict,img):
    img = img[:,:,0] if len(img.shape) == 3 else img
    img_out = np.zeros(img.shape + (3,))
#變成RGB空間，因為其他顏色只能再RGB空間才會顯示
    for i in range(num_class):
        img_out[img == i,:] = color_dict[i]
#為不同類別塗上不同的顏色，color_dict[i]是與類別數有關的顏色，img_out[img == i,:]是img_out在img中等於i類的位置上的點
    return img_out / 255
 
#上面函數是給出測試后的輸出之后，為輸出塗上不同的顏色，多類情況下才起作用，兩類的話無用
 
def saveResult(save_path,npyfile,flag_multi_class = False,num_class = 2):
    for i,item in enumerate(npyfile):
        img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]
#多類的話就圖成彩色，非多類（兩類）的話就是黑白色
        io.imsave(os.path.join(save_path,"%d_predict.png"%i),img)19
#下面是model.py：

 

import numpy as np 
import os
import skimage.io as io
import skimage.transform as trans
import numpy as np
from keras.models import *
from keras.layers import *
from keras.optimizers import *
from keras.callbacks import ModelCheckpoint, LearningRateScheduler
from keras import backend as keras
 
 
def unet(pretrained_weights = None,input_size = (256,256,1)):
    inputs = Input(input_size)
    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)
    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)
    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)
    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)
    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)
    drop4 = Dropout(0.5)(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
 
    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)
    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)
    drop5 = Dropout(0.5)(conv5)
 
    up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))#上采樣之后再進行卷積，相當於轉置卷積操作！
    merge6 = concatenate([drop4,up6],axis=3)
    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)
 
    up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))
    merge7 = concatenate([conv3,up7],axis = 3)
    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
 
    up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))
    merge8 = concatenate([conv2,up8],axis = 3)
    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)
    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)
 
    up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))
    merge9 = concatenate([conv1,up9],axis = 3)
    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)
    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
    conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
    conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)#我懷疑這個sigmoid激活函數是多余的，因為在后面的loss中用到的就是二進制交叉熵，包含了sigmoid
 
    model = Model(input = inputs, output = conv10)
 
    model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])#模型執行之前必須要編譯https://keras-cn.readthedocs.io/en/latest/getting_started/sequential_model/
    #利用二進制交叉熵，也就是sigmoid交叉熵，metrics一般選用准確率，它會使准確率往高處發展
    #model.summary()
 
    if(pretrained_weights):
     model.load_weights(pretrained_weights)
 
    return model