在本节开始,就是正式搭建Unet模型了。在本次搭建模型的过程中,将采用一定的封装来搭建网络,因为U-net模型有个显著的特点,就是很多重复的步骤,比如下采样、两层卷积——>下采样、两层卷积......因此可以做一个封装让程序更加简化。

前置知识:

U-NET模型结构

在创建模型的时候,对上采样和下采样进行了封装,分别为类Downsample和类Upsample

在下采样层中,设置了一个标识符is_pool,这是因为在上图画框的地方能用到的只有两层卷积,通过标识符来判断是否下采样,可以更加充分地利用自定义的类。

此外,还自定义可IOU,因为预测的结果是34个张量值,而标签是具体的34个分类(int类型值),因此需要进行转换:y_pred = tf.argmax(y_pred,axis=-1)

为了在训练中引入IOU,因此使用自定义训练。

完整代码:

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

# E:\citydata
# 获取train文件夹里的所有城市及所有图片
img = glob.glob('E:/citydata/leftImg8Bit/train/*/*.png')
# 获取标签
label = glob.glob('E:/citydata/gtFine/train/*/*_gtFine_labelIds.png')

# 把顺序打乱,但是保持label和img仍然是一一对应的
index = np.random.permutation(len(img))
img = np.array(img)[index]
label = np.array(label)[index]

# 测试集
img_val = glob.glob('E:/citydata/leftImg8Bit/val/*/*.png')
label_val = glob.glob('E:/citydata/gtFine/val/*/*_gtFine_labelIds.png')

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

# 图像的加载
# img是png彩色图像
def read_png(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img,channels=3)
    return img

# 标签图片,通道为1
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)  
    concat_img = tf.image.resize(concat_img,(280,280),
                                 method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    crop_img = tf.image.random_crop(concat_img,[256,256,4])
    # 使用高级切片的方式,获得分开的原图和标签
    return crop_img[:,:,:3],crop_img[:,:,3:]

def normal(img, mask):
    #[-1,1]
    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:
        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 = 8
BUFFER_SIZE = 50

train_count = len(label)
test_count = len(label_val)

step_per_epoch = train_count//BATCH_SIZE
val_step = test_count//BATCH_SIZE

AUTO = tf.data.experimental.AUTOTUNE

dataset_train = dataset_train.map(load_image_train,num_parallel_calls=AUTO)
dataset_test = dataset_test.map(load_image_test,num_parallel_calls=AUTO)

dataset_train = dataset_train.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
dataset_test = dataset_test.batch(BATCH_SIZE)

# 定义下采样层
class Downsample(tf.keras.layers.Layer):
    def __init__(self,units) :
        super(Downsample,self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(units,
                                            kernel_size=3,
                                            padding='same'
                                           )
        self.conv2 = tf.keras.layers.Conv2D(units,
                                            kernel_size=3,
                                            padding='same'
                                           )
        self.pool = tf.keras.layers.MaxPooling2D()
    # 默认是进行下采样的,但是有的地方不需要,因此设置is_pool个节点
    def call(self,x,is_pool=True):
        if is_pool:  #如果is_pool的值是True,进行下采样
            x = self.pool(x)
            
        x = self.conv1(x)
        x = tf.nn.relu(x)
        
        x = self.conv2(x)
        x = tf.nn.relu(x)       
        
        return x
    
# 定义上采样层
class Upsample(tf.keras.layers.Layer):
    def __init__(self,units) :
        super(Upsample,self).__init__()    
        self.conv1 = tf.keras.layers.Conv2D(units,
                                            kernel_size=3,
                                            padding='same'
                                           )
        self.conv2 = tf.keras.layers.Conv2D(units,
                                            kernel_size=3,
                                            padding='same'
                                           )    
        # 反卷积的单元数是之前的一般
        self.deconv = tf.keras.layers.Conv2DTranspose(units//2,
                                                      kernel_size=3,
                                                      strides=2,
                                                      padding='same'
                                                     )
    def call(self,x,is_pool=True):
        x = self.conv1(x)
        x = tf.nn.relu(x)
        
        x = self.conv2(x)
        x = tf.nn.relu(x)       
        
        x = self.deconv(x)
        x = tf.nn.relu(x)
        
        return x                

class Unet_model(tf.keras.Model):
    def __init__(self) :
        super(Unet_model,self).__init__()        
        # 第一层就是下采样
        self.down1 = Downsample(64)
        self.down2 = Downsample(128)
        self.down3 = Downsample(256)
        self.down4 = Downsample(512)
        self.down5 = Downsample(1024)
        
        self.up = tf.keras.layers.Conv2DTranspose(512,
                                                 kernel_size = 2,
                                                 strides = 2,
                                                 padding='same'
                                                 )
        self.up1 = Upsample(512)      
        self.up2 = Upsample(256) 
        self.up3 = Upsample(128) 
        
        self.conv_last = Downsample(64)
        
        self.last = tf.keras.layers.Conv2D(34,
                                          kernel_size=1,
                                          padding='same')
    def call(self,x):
        x1 = self.down1(x,is_pool=False)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.down5(x4)
        
        x5 = self.up(x5)
        
        x5 = tf.concat([x4,x5], axis=-1)
        x5 = self.up1(x5)
        x5 = tf.concat([x3,x5], axis=-1)
        x5 = self.up2(x5)
        x5 = tf.concat([x2,x5], axis=-1)
        x5 = self.up3(x5)
        
        x5 = tf.concat([x1,x5], axis=-1)
        
        x5 = self.conv_last(x5,is_pool=False)
        
        x5 = self.last(x5)
        
        return x5
    
model = Unet_model()

opt = tf.keras.optimizers.Adam(0.0001)

# labels 0,1,2,3
# 前面没有激活,from_logits设置True
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) 

# 把预测的值转换为分类的label数值
class MeanIOU(tf.keras.metrics.MeanIoU):
    def __call__(self,y_true,y_pred):
        y_pred = tf.argmax(y_pred,axis=-1)
        return super().__call__(y_true,y_pred)

train_loss = tf.keras.metrics.Mean(name="train_loss")
train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name="train_acc")
train_iou = MeanIOU(34, name="train_iou")

test_loss = tf.keras.metrics.Mean(name="test_loss")
test_acc = tf.keras.metrics.SparseCategoricalAccuracy(name="test_acc")
test_iou = MeanIOU(34, name="test_iou")

@tf.function
def train_step(images,labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_fn(labels,predictions)
    gradients = tape.gradient(loss,model.trainable_variables)
    opt.apply_gradients(zip(gradients,model.trainable_variables))
    
    train_loss(loss)
    train_acc(labels, predictions)
    train_iou(labels,predictions)
    
@tf.function
def test_step(images,labels):
    predictions = model(images)
    loss = loss_fn(labels,predictions)
    
    test_loss(loss)
    test_acc(labels, predictions)
    test_iou(labels,predictions)
             
EPOCHS = 5
for epoch in range(EPOCHS):
    # 在下一个epoch开始时,重置评估指标
    train_loss.reset_states()
    train_acc.reset_states()
    train_iou.reset_states()
    
    test_acc.reset_states()
    test_loss.reset_states()
    test_iou.reset_states()
    
    for images, labels in dataset_train:
        train_step(images, labels)

    for test_images, test_labels in dataset_test:
        test_step(test_images, test_labels)

    template = 'Epoch {:.3f}, Loss: {:.3f}, Accuracy: {:.3f}, \
                IOU: {:.3f}, Test Loss: {:.3f}, \
                Test Accuracy: {:.3f}, Test IOU: {:.3f}'
    print (template.format(epoch+1,
                           train_loss.result(),
                           train_acc.result()*100,
                           train_iou.result(),
                           test_loss.result(),
                           test_acc.result()*100,
                           test_iou.result() 
                           ))

如果你的电脑配置比较高,可以多跑些epoch。

训练的结果


版权声明 ▶ 本网站名称:陶小桃Blog
▶ 本文链接:https://www.52txr.cn/2023/unet5.html
▶ 本网站的文章部分内容可能来源于网络,仅供大家学习与参考,如有侵权,请联系站长进行核实删除。
▶ 转载本站文章需要遵守:商业转载请联系站长,非商业转载请注明出处并附带原文链接!!
▶ 站长邮箱 [email protected][email protected] ,如不方便留言可邮件联系。

小陶的个人微信公众号

学累了就来张美女照片养养眼吧,身体是革命的本钱,要多多休息哦~ 随机美女图片

最后修改:2023 年 03 月 23 日
如果觉得我的文章对你有用,请随意赞赏!