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

获取img、label路径

1
2
img = glob.glob('./images/leftImg8bit/train/*/*.png')#tf.io.glob.glob
label = glob.glob('./images/gtFine/train/*/*_gtFine_labelIds.png')

获取img_val、label_val路径

1
2
img_val = glob.glob('./images/leftImg8bit/val/*/*.png')
label_val = glob.glob('./images/gtFine/val/*/*_gtFine_labelIds.png')

对训练路径进行乱序（减轻shuffle）

1
2
3
index = np.random.permutation(len(img))
img = np.array(img)[index]
label = np.array(label)[index]

创建输入管道

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

创建数据增强函数

tf.random.uniform()#返回从0~1之间的均匀分布
1.tf.image.flip_left_right() #随机翻转
2.随机裁剪 2张图片合并叠加在一起剪完拆开

1
2
3
4
5
6
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:]

调整大小的图像将失真,如果他们的原始纵横比与size不一样.为了避免扭曲,请参阅tf.image.resize_image_with_crop_or_pad.

method 可以是下列之一：

ResizeMethod.BILINEAR：双线性插值.
ResizeMethod.NEAREST_NEIGHBOR：最近的邻居插值.
ResizeMethod.BICUBIC：双三次插值.
ResizeMethod.AREA：区域插值.
如果method是ResizeMethod.NEAREST_NEIGHBOR,则返回值与images具有相同的类型.它也将与images具有相同的类型,如果images的大小可以静态地确定为与size相同,因为在这种情况下返回images.否则,返回值有类型float32.

参数：

images：形状为[batch, height, width, channels]的4-D张量或形状为[height, width, channels]的3-D张量.
size：2个元素(new_height, new_width)的1维int32张量,表示图像的新大小.
method：ResizeMethod,默认为ResizeMethod.BILINEAR.
align_corners：布尔型,如果为True,则输入和输出张量的4个拐角像素的中心对齐,并且保留角落像素处的值；默认为False.
可能引发的异常：

ValueError：如果images的形状与此函数的形状参数不兼容.
ValueError：如果size有无效的形状或类型.
ValueError：如果指定了不支持的调整大小方法.
返回：

如果images是四维,则返回一个形状为[batch, new_height, new_width, channels]的四维浮动张量；如果images是三维,则返回一个形状为[new_height, new_width, channels]的三维浮动张量.

创建图片读取函数

1
2
3
4
def read_png(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img,channels=3)
    return img

创建图片对应label图片读取函数

1
2
3
4
def read_png_label(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img,channels=1)
    return img

创建归一化函数

1
2
3
4
def normal(img,mask):#归一化 
    img = tf.cast(img,tf.float32)/127.5-1
    mask = tf.cast(mask,tf.int32)
    return img,mask

创建训练图片读取函数

1
2
3
4
5
6
7
8
9
10
11
12
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~1之间的均匀分布  '()'表示生成常量
        img = tf.image.flip_left_right(img)
        mask = tf.image.flip_left_right(mask)
        
    img,mask = normal(img,mask)
    return img,mask

创建测试图片读取函数

1
2
3
4
5
6
7
8
9
10
def load_image_val(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

设置AUTOTUNE

1
AUTOTUNE = tf.data.experimental.AUTOTUNE

将函数映射到Dataset

1
2
dataset_train = dataset_train.map(load_image_train,num_parallel_calls=AUTOTUNE)
dataset_val = dataset_val.map(load_image_val,num_parallel_calls=AUTOTUNE)

观察数据

1
2
3
4
5
for i,m in dataset_train.take(1):
    plt.subplot(1,2,1)
    plt.imshow((i.numpy()+1)/2)
    plt.subplot(1,2,2)
    plt.imshow(np.squeeze(m.numpy()))#squeeze 函数：从数组的形状中删除单维度条目，即把shape中为1的维度去掉

设置BATCH_SIZE,BUFFER_SIZE,step_per_epoch，val_step

1
2
3
4
5
6
7
BATCH_SIZE = 6
BUFFER_SIZE = 300
step_per_epoch = len(img)//BATCH_SIZE
val_step = val_count//BATCH_SIZE

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

定义unet模型

定义下采样层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
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()
    def call(self, x, is_pool=True):
        if is_pool:
            x = self.pool(x)
        x = self.conv1(x)
        x = tf.nn.relu(x)
        x = self.conv2(x)
        x = tf.nn.relu(x)
        return x

定义上采样层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class Upsample(tf.keras.layers.Layer):
    def __init__(self, unit):
        super(Upsample, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(unit, kernel_size=3, padding='same')
        self.conv2 = tf.keras.layers.Conv2D(unit, kernel_size=3, padding='same')
        self.deconv = tf.keras.layers.Conv2DTranspose(unit//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

定义Unet模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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.up1 = Upsample(512)
        self.up2 = Upsample(256)
        self.up3 = Upsample(128)
        
        self.conv_last = Downsample(64)
        
        self.up = tf.keras.layers.Conv2DTranspose(512,
                                               kernel_size=3,
                                               strides = 2,
                                               padding='same')
        
        self.last = tf.keras.layers.Conv2D(34,
                                        kernel_size=1,
                                        padding='same')

    def call(self, x):
        x1 = self.down1(x, is_pool=False)         # 256*256*64
        x2 = self.down2(x1)                       # 128*128*128
        x3 = self.down3(x2)                       # 64*64*256
        x4 = self.down4(x3)                       # 32*32*512
        x5 = self.down5(x4)                       # 16*16*1024
        
        x5 = self.up(x5)                          # 32*32*512
        
        x5 = tf.concat([x4,x5], axis=-1)          # 32*32*1024
        x5 = self.up1(x5)                         # 64*64*256
        x5 = tf.concat([x3,x5], axis=-1)          # 64*64*512   
        x5 = self.up2(x5)                         # 128*128*128
        x5 = tf.concat([x2,x5], axis=-1)          # 128*128*256
        x5 = self.up3(x5)                         # 256*256*64
        x5 = tf.concat([x1,x5], axis=-1)          # 256*256*128
        
        x5 = self.conv_last(x5, is_pool=False)    # 256*256*64
        
        x5 = self.last(x5)                        # 256*256*3
        return x5

设置IOU交并比

1
2
3
4
class MeanIOU(tf.keras.metrics.MeanIoU):#tf.keras.metrics.MeanIoU平均交并比  因为y_true 对于每一个像素长度34的张量 y_pred类别不统一要处理
    def __call__(self, y_true, y_pred, sample_weight=None):#y_true 对于每一个像素长度34的张量 y_pred类别
        y_pred = tf.argmax(y_pred, axis=-1)
        return super().__call__(y_true, y_pred, sample_weight=sample_weight)

创建Unet模型对象

1
model = Unet_model()

设置optimizer

1
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

设置计算accuracy、loss、iou函数

1
2
3
4
5
6
7
8
9
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)#数字编码

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

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

设置train_step

1
2
3
4
5
6
7
8
9
10
11
@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)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)
    train_iou(labels, predictions)

设置test_step

1
2
3
4
5
6
7
8
@tf.function
def test_step(images,labels):
    predictions = model(images)
    t_loss = loss_fn(labels,predictions)
    
    test_loss(t_loss)
    test_accuracy(labels,predictions)
    test_iou(labels,predictions)

开始训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
EPOCHS = 60
for epoch in range(EPOCHS):
    #在下一个epoch开始时，重置评估指标
    train_loss.reset_states()
    train_accuracy.reset_states()
    train_iou.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()
    test_iou.reset_states()
    
    for images, labels in dataset_train:
        train_step(images, labels)
        
    for test_images,test_labels in dataset_val:
        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_accuracy.result()*100,
                          train_iou.result(),
                          test_loss.result(),
                          test_accuracy.result()*100,
                          test_iou.result()
                         ))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
Epoch1.000,Loss1.680,Accuracy52.866,                IOU0.058,Test Loss1.303,                Test Accuracy64.521,Test IOU0.086
Epoch2.000,Loss1.035,Accuracy70.013,                IOU0.110,Test Loss1.151,                Test Accuracy69.294,Test IOU0.110
Epoch3.000,Loss0.903,Accuracy74.234,                IOU0.130,Test Loss1.089,                Test Accuracy69.691,Test IOU0.119
Epoch4.000,Loss0.798,Accuracy77.655,                IOU0.150,Test Loss0.895,                Test Accuracy76.077,Test IOU0.156
Epoch5.000,Loss0.719,Accuracy80.269,                IOU0.170,Test Loss0.861,                Test Accuracy77.156,Test IOU0.159
Epoch6.000,Loss0.666,Accuracy81.814,                IOU0.188,Test Loss0.859,                Test Accuracy77.956,Test IOU0.178
Epoch7.000,Loss0.632,Accuracy82.718,                IOU0.199,Test Loss0.818,                Test Accuracy79.519,Test IOU0.186
Epoch8.000,Loss0.611,Accuracy83.235,                IOU0.207,Test Loss0.807,                Test Accuracy79.951,Test IOU0.192
Epoch9.000,Loss0.589,Accuracy83.748,                IOU0.216,Test Loss0.732,                Test Accuracy81.035,Test IOU0.207
Epoch10.000,Loss0.567,Accuracy84.253,                IOU0.227,Test Loss0.734,                Test Accuracy80.651,Test IOU0.208
Epoch11.000,Loss0.554,Accuracy84.574,                IOU0.234,Test Loss0.732,                Test Accuracy81.601,Test IOU0.214
Epoch12.000,Loss0.536,Accuracy84.950,                IOU0.243,Test Loss0.737,                Test Accuracy81.746,Test IOU0.215
Epoch13.000,Loss0.534,Accuracy84.945,                IOU0.243,Test Loss0.685,                Test Accuracy82.114,Test IOU0.229
Epoch14.000,Loss0.515,Accuracy85.454,                IOU0.256,Test Loss0.698,                Test Accuracy81.972,Test IOU0.226
Epoch15.000,Loss0.507,Accuracy85.640,                IOU0.262,Test Loss0.724,                Test Accuracy81.256,Test IOU0.226
Epoch16.000,Loss0.490,Accuracy86.033,                IOU0.270,Test Loss0.665,                Test Accuracy82.712,Test IOU0.242
Epoch17.000,Loss0.485,Accuracy86.154,                IOU0.276,Test Loss0.699,                Test Accuracy82.372,Test IOU0.240
Epoch18.000,Loss0.476,Accuracy86.337,                IOU0.280,Test Loss0.688,                Test Accuracy81.953,Test IOU0.239
Epoch19.000,Loss0.471,Accuracy86.506,                IOU0.284,Test Loss0.668,                Test Accuracy82.952,Test IOU0.252
Epoch20.000,Loss0.465,Accuracy86.629,                IOU0.290,Test Loss0.685,                Test Accuracy82.774,Test IOU0.235
Epoch21.000,Loss0.455,Accuracy86.854,                IOU0.296,Test Loss0.687,                Test Accuracy82.975,Test IOU0.251
Epoch22.000,Loss0.447,Accuracy87.007,                IOU0.300,Test Loss0.644,                Test Accuracy83.284,Test IOU0.257
Epoch23.000,Loss0.441,Accuracy87.169,                IOU0.306,Test Loss0.687,                Test Accuracy82.993,Test IOU0.256
Epoch24.000,Loss0.445,Accuracy87.117,                IOU0.308,Test Loss0.659,                Test Accuracy82.803,Test IOU0.258
Epoch25.000,Loss0.433,Accuracy87.353,                IOU0.313,Test Loss0.657,                Test Accuracy83.123,Test IOU0.256
Epoch26.000,Loss0.424,Accuracy87.637,                IOU0.324,Test Loss0.726,                Test Accuracy82.877,Test IOU0.262
Epoch27.000,Loss0.424,Accuracy87.594,                IOU0.323,Test Loss0.701,                Test Accuracy83.324,Test IOU0.262
Epoch28.000,Loss0.420,Accuracy87.694,                IOU0.326,Test Loss0.665,                Test Accuracy83.373,Test IOU0.261
Epoch29.000,Loss0.417,Accuracy87.782,                IOU0.328,Test Loss0.658,                Test Accuracy83.086,Test IOU0.265
Epoch30.000,Loss0.410,Accuracy87.926,                IOU0.333,Test Loss0.652,                Test Accuracy83.806,Test IOU0.279
Epoch31.000,Loss0.408,Accuracy87.994,                IOU0.335,Test Loss0.688,                Test Accuracy83.246,Test IOU0.266
Epoch32.000,Loss0.403,Accuracy88.120,                IOU0.336,Test Loss0.662,                Test Accuracy82.924,Test IOU0.268
Epoch33.000,Loss0.404,Accuracy88.081,                IOU0.339,Test Loss0.725,                Test Accuracy83.613,Test IOU0.272
Epoch34.000,Loss0.398,Accuracy88.230,                IOU0.346,Test Loss0.681,                Test Accuracy83.905,Test IOU0.281
Epoch35.000,Loss0.391,Accuracy88.434,                IOU0.351,Test Loss0.679,                Test Accuracy84.135,Test IOU0.273
Epoch36.000,Loss0.389,Accuracy88.485,                IOU0.354,Test Loss0.704,                Test Accuracy83.945,Test IOU0.265
Epoch37.000,Loss0.393,Accuracy88.393,                IOU0.350,Test Loss0.646,                Test Accuracy84.268,Test IOU0.280
Epoch38.000,Loss0.381,Accuracy88.676,                IOU0.360,Test Loss0.686,                Test Accuracy84.167,Test IOU0.279
Epoch39.000,Loss0.386,Accuracy88.535,                IOU0.361,Test Loss0.671,                Test Accuracy84.283,Test IOU0.278
Epoch40.000,Loss0.384,Accuracy88.649,                IOU0.360,Test Loss0.652,                Test Accuracy83.647,Test IOU0.285
Epoch41.000,Loss0.375,Accuracy88.910,                IOU0.369,Test Loss0.645,                Test Accuracy83.704,Test IOU0.285
Epoch42.000,Loss0.376,Accuracy88.835,                IOU0.370,Test Loss0.648,                Test Accuracy83.660,Test IOU0.282
Epoch43.000,Loss0.384,Accuracy88.665,                IOU0.364,Test Loss0.673,                Test Accuracy84.337,Test IOU0.284
Epoch44.000,Loss0.368,Accuracy89.020,                IOU0.381,Test Loss0.681,                Test Accuracy84.040,Test IOU0.283
Epoch45.000,Loss0.362,Accuracy89.195,                IOU0.383,Test Loss0.692,                Test Accuracy83.273,Test IOU0.279
Epoch46.000,Loss0.371,Accuracy88.986,                IOU0.377,Test Loss0.663,                Test Accuracy84.442,Test IOU0.294
Epoch47.000,Loss0.361,Accuracy89.252,                IOU0.390,Test Loss0.627,                Test Accuracy84.221,Test IOU0.292
Epoch48.000,Loss0.368,Accuracy89.076,                IOU0.378,Test Loss0.701,                Test Accuracy83.599,Test IOU0.284
Epoch49.000,Loss0.366,Accuracy89.096,                IOU0.385,Test Loss0.695,                Test Accuracy84.051,Test IOU0.286
Epoch50.000,Loss0.372,Accuracy88.946,                IOU0.383,Test Loss0.683,                Test Accuracy84.060,Test IOU0.291
Epoch51.000,Loss0.352,Accuracy89.501,                IOU0.403,Test Loss0.683,                Test Accuracy84.378,Test IOU0.292
Epoch52.000,Loss0.360,Accuracy89.266,                IOU0.395,Test Loss0.688,                Test Accuracy84.026,Test IOU0.292
Epoch53.000,Loss0.346,Accuracy89.610,                IOU0.405,Test Loss0.712,                Test Accuracy84.051,Test IOU0.288
Epoch54.000,Loss0.354,Accuracy89.423,                IOU0.402,Test Loss0.673,                Test Accuracy84.827,Test IOU0.307
Epoch55.000,Loss0.355,Accuracy89.481,                IOU0.400,Test Loss0.650,                Test Accuracy84.651,Test IOU0.294
Epoch56.000,Loss0.349,Accuracy89.605,                IOU0.407,Test Loss0.699,                Test Accuracy84.443,Test IOU0.292
Epoch57.000,Loss0.353,Accuracy89.483,                IOU0.403,Test Loss0.688,                Test Accuracy84.417,Test IOU0.291
Epoch58.000,Loss0.350,Accuracy89.564,                IOU0.411,Test Loss0.693,                Test Accuracy84.770,Test IOU0.298
Epoch59.000,Loss0.344,Accuracy89.700,                IOU0.410,Test Loss0.684,                Test Accuracy84.292,Test IOU0.291
Epoch60.000,Loss0.340,Accuracy89.859,                IOU0.417,Test Loss0.649,                Test Accuracy84.367,Test IOU0.302

模型权重保存

1
model.save_weights('city_weights.h5')

模型图片预测

1
2
3
4
5
6
7
8
test_path = 'test.png'
test_png = read_png(test_path)
test_png = tf.image.resize(test_png,(256,256),
                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
test_png = tf.cast(test_png,tf.float32)/255
test_png = np.expand_dims(test_png,0)#扩张维度 
test_pred = model.predict(test_png)
test_p = np.argmax(test_pred,axis=-1)
1
2
3
4
plt.subplot(1,2,1)
plt.imshow(np.squeeze(test_png.numpy()))
plt.subplot(1,2,2)
plt.imshow(np.squeeze(test_p))#squeeze 函数：从数组的形状中删除单维度条目，即把shape中为1的维度去掉

==np.squeeze(m.numpy()))#squeeze 函数：从数组的形状中删除单维度条目，即把shape中为1的维度去掉==