深度學習-Tensorflow2.2-圖像處理{10}-UNET圖像語義分割模型-24

UNET圖像語義分割模型簡介

代碼

import tensorflow as tf
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import glob
import os

# 顯存自適應分配
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu,True)

gpu_ok = tf.test.is_gpu_available()
print("tf version:", tf.__version__)
print("use GPU", gpu_ok) # 判斷是否使用gpu進行訓練

獲取訓練數據及目標值

# 獲取train文件下所有文件中所有png的圖片
img = glob.glob("G:/BaiduNetdiskDownload/cityscapes/leftImg8bit/train/*/*.png") 

train_count = len(img)
img[:5],train_count

# 獲取gtFine/train文件下所有文件中所有_gtFine_labelIds.png的圖片
label = glob.glob("G:/BaiduNetdiskDownload/cityscapes/gtFine/train/*/*_gtFine_labelIds.png") 

index = np.random.permutation(len(img)) # 創建一個隨即種子，保障image和label 隨機后還是一一對應的

img = np.array(img)[index] # 對訓練集圖片進行亂序
label = np.array(label)[index]

獲取測試數據

# 獲取val文件下所有文件中所有png的圖片
img_val = glob.glob("G:/BaiduNetdiskDownload/cityscapes/leftImg8bit/val/*/*.png") 

# 獲取gtFine/val文件下所有文件中所有_gtFine_labelIds.png的圖片
label_val = glob.glob("G:/BaiduNetdiskDownload/cityscapes/gtFine/val/*/*_gtFine_labelIds.png") 

test_count = len(img_val)
img_val[:5],test_count,label_val[:5],len(label_val)

創建數據集

dataset_train = tf.data.Dataset.from_tensor_slices((img,label))
dataset_val = tf.data.Dataset.from_tensor_slices((img_val,label_val))

# 創建png的解碼函數
def read_png(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img,channels=3)
    return img

# 創建png的解碼函數
def read_png_label(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img,channels=1)
    return img

# 數據增強
def crop_img(img,mask):
    concat_img = tf.concat([img,mask],axis=-1) # 把image和label合並在一起 axis = -1，表示最后一個維度
    concat_img = tf.image.resize(concat_img,(280,280), # 修改大小為280*280
                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)#使用最近鄰插值調整images為size
    crop_img = tf.image.random_crop(concat_img,[256,256,4]) # 隨機裁剪
    return crop_img[ :, :, :3],crop_img[ :, :, 3:] # 高維切片(第一，第二維度全要，第三個維度的前3是image，最后一個維度就是label)

def normal(img,mask):
    img = tf.cast(img,tf.float32)/127.5-1
    mask = tf.cast(mask,tf.int32)
    return img,mask

# 組裝
def load_image_train(img_path,mask_path):
    img = read_png(img_path)
    mask = read_png_label(mask_path) # 獲取路徑
    
    img,mask = crop_img(img,mask) # 調用隨機裁剪函數對圖片進行裁剪
    
    if tf.random.uniform(())>0.5: # 從均勻分布中返回隨機值 如果大於0.5就執行下面的隨機翻轉
        img = tf.image.flip_left_right(img)
        mask = tf.image.flip_left_right(mask)
    img,mask = normal(img,mask) # 調用歸一化函數
    return img,mask

# 組裝
def load_image_test(img_path,mask_path):
    img = read_png(img_path)
    mask = read_png_label(mask_path)
    
    img = tf.image.resize(img,(256,256))
    mask = tf.image.resize(mask,(256,256))
    
    img,mask = normal(img,mask)
    
    return img,mask

BATCH_SIZE = 32
BUFFER_SIZE = 300
step_per_epoch = train_count//BATCH_SIZE
val_step = test_count//BATCH_SIZE

auto = tf.data.experimental.AUTOTUNE # 根據cpu使用情況自動規划線程讀取圖片

# 創建輸入管道
dataset_train = dataset_train.map(load_image_train,num_parallel_calls=auto)
dataset_val = dataset_val.map(load_image_test,num_parallel_calls=auto)

dataset_train = dataset_train.cache().repeat().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).prefetch(auto)
dataset_val = dataset_val.cache().batch(BATCH_SIZE)

定義unet模型

def create_model():
    ## unet網絡結構下采樣部分
    # 輸入層 第一部分
    inputs = tf.keras.layers.Input(shape = (256,256,3))
    
    x = tf.keras.layers.Conv2D(64,3,padding="same",activation="relu")(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Conv2D(64,3,padding="same",activation="relu")(x)
    x = tf.keras.layers.BatchNormalization()(x)  # 256*256*64
    # 下采樣
    x1 = tf.keras.layers.MaxPooling2D(padding="same")(x) # 128*128*64
    
    # 卷積 第二部分
    x1 = tf.keras.layers.Conv2D(128,3,padding="same",activation="relu")(x1)
    x1 = tf.keras.layers.BatchNormalization()(x1)
    x1 = tf.keras.layers.Conv2D(128,3,padding="same",activation="relu")(x1)
    x1 = tf.keras.layers.BatchNormalization()(x1)  # 128*128*128
    # 下采樣 
    x2 = tf.keras.layers.MaxPooling2D(padding="same")(x1) # 64*64*128
        
    # 卷積 第三部分
    x2 = tf.keras.layers.Conv2D(256,3,padding="same",activation="relu")(x2)
    x2 = tf.keras.layers.BatchNormalization()(x2)
    x2 = tf.keras.layers.Conv2D(256,3,padding="same",activation="relu")(x2)
    x2 = tf.keras.layers.BatchNormalization()(x2)  # 64*64*256
    # 下采樣
    x3 = tf.keras.layers.MaxPooling2D(padding="same")(x2) # 32*32*256
    
    # 卷積 第四部分
    x3 = tf.keras.layers.Conv2D(512,3,padding="same",activation="relu")(x3)
    x3 = tf.keras.layers.BatchNormalization()(x3)
    x3 = tf.keras.layers.Conv2D(512,3,padding="same",activation="relu")(x3)
    x3 = tf.keras.layers.BatchNormalization()(x3)  # 32*32*512
    # 下采樣
    x4 = tf.keras.layers.MaxPooling2D(padding="same")(x3) # 16*16*512
    # 卷積 第五部分
    x4 = tf.keras.layers.Conv2D(1024,3,padding="same",activation="relu")(x4)
    x4 = tf.keras.layers.BatchNormalization()(x4)
    x4 = tf.keras.layers.Conv2D(1024,3,padding="same",activation="relu")(x4)
    x4 = tf.keras.layers.BatchNormalization()(x4)  # 16*16*1024
    
    ## unet網絡結構上采樣部分
    
    # 反卷積 第一部分 512個卷積核 卷積核大小2*2 跨度2 填充方式same 激活relu
    x5 = tf.keras.layers.Conv2DTranspose(512,2,strides=2,
                                         padding="same",
                                         activation="relu")(x4)#32*32*512
    x5 = tf.keras.layers.BatchNormalization()(x5)
    x6 = tf.concat([x3,x5],axis=-1)#合並 32*32*1024
    # 卷積
    x6 = tf.keras.layers.Conv2D(512,3,padding="same",activation="relu")(x6)
    x6 = tf.keras.layers.BatchNormalization()(x6)
    x6 = tf.keras.layers.Conv2D(512,3,padding="same",activation="relu")(x6)
    x6 = tf.keras.layers.BatchNormalization()(x6)  # 32*32*512
    
    # 反卷積 第二部分
    x7 = tf.keras.layers.Conv2DTranspose(256,2,strides=2,
                                         padding="same",
                                         activation="relu")(x6)#64*64*256
    x7 = tf.keras.layers.BatchNormalization()(x7)
    x8 = tf.concat([x2,x7],axis=-1)#合並 64*64*512
    # 卷積
    x8 = tf.keras.layers.Conv2D(256,3,padding="same",activation="relu")(x8)
    x8 = tf.keras.layers.BatchNormalization()(x8)
    x8 = tf.keras.layers.Conv2D(256,3,padding="same",activation="relu")(x8)
    x8 = tf.keras.layers.BatchNormalization()(x8)  # #64*64*256
    
    # 反卷積 第三部分
    x9 = tf.keras.layers.Conv2DTranspose(128,2,strides=2,
                                         padding="same",
                                         activation="relu")(x8)# 128*128*128
    x9 = tf.keras.layers.BatchNormalization()(x9)
    x10 = tf.concat([x1,x9],axis=-1)#合並 128*128*256
    # 卷積
    x10 = tf.keras.layers.Conv2D(128,3,padding="same",activation="relu")(x10)
    x10 = tf.keras.layers.BatchNormalization()(x10)
    x10 = tf.keras.layers.Conv2D(128,3,padding="same",activation="relu")(x10)
    x10 = tf.keras.layers.BatchNormalization()(x10)  # 128*128*128
    
    # 反卷積 第四部分
    x11 = tf.keras.layers.Conv2DTranspose(64,2,strides=2,
                                         padding="same",
                                         activation="relu")(x10)# 256*256*64
    x11 = tf.keras.layers.BatchNormalization()(x11)
    x12 = tf.concat([x,x11],axis=-1)#合並 256*256*128
    # 卷積
    x12 = tf.keras.layers.Conv2D(64,3,padding="same",activation="relu")(x12)
    x12 = tf.keras.layers.BatchNormalization()(x12)
    x12 = tf.keras.layers.Conv2D(64,3,padding="same",activation="relu")(x12)
    x12 = tf.keras.layers.BatchNormalization()(x12)  # 256*256*64
    
    # 輸出層 第五部分
    output =tf.keras.layers.Conv2D(34,1,padding="same",activation="softmax")(x12)# 256*256*34
    
    return tf.keras.Model(inputs=inputs,outputs=output)
    

model = create_model()

tf.keras.utils.plot_model(model) # 繪制模型圖

# tf.keras.metrics.MeanIoU(num_classes=34) # 根據獨熱編碼進行計算
# 我們是順序編碼 需要更改類
class MeanIou(tf.keras.metrics.MeanIoU): # 繼承這個類 
    def __call__(self,y_true,y_pred,sample_weight=None): 
        y_pred = tf.argmax(u_pred,axis=-1) 
        return super().__call__(y_true,y_pred,sample_weight=sample_weight)

# 編譯模型
model.compile(optimizer="adam",
              loss="sparse_categorical_crossentropy",
              metrics=["acc",MeanIou(num_classes=34)]
             )

# 訓練
history = model.fit(dataset_train,
                    epochs=60,
                    steps_per_epoch=step_per_epoch,
                    validation_steps=val_step,
                    validation_data=dataset_val
                   )

列子