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
                   )

列子