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

查看Tensorflow版本

1
tf.__version__

仅在需要时申请显存空间

程序初始运行时消耗 很少的显存，随着程序的运行而动态申请显存；GPU 的显存使用策略设置为 “仅在需要时申请显存空间”)

1
2
3
gpus = tf.config.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(device=gpu, enable=True)

glob读取路径

训练数据

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

训练数据乱序

1
2
3
4
index = np.random.permutation(len(img))

img = np.array(img)[index]
label = np.array(label)[index]

测试数据

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

Dataset

训练数据Dataset

1
dataset_train = tf.data.Dataset.from_tensor_slices((img, label))

测试数据Dataset

1
dataset_val = tf.data.Dataset.from_tensor_slices((img_val, label_val))

图片读取函数

1
2
3
4
def read_png(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img, channels=3)
    return img
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.随机翻转 img = tf.image.flip_left_right()

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:]

归一化函数

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
13
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

测试图片加载

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

设置BACTH_SIZE、BUFFER_SIZE、step_EPOCH

1
2
3
4
BCTCH_SIZE = 8
BUFFER_SIZE = 300
step_per_epoch = train_count//BCTCH_SIZE
val_step = val_count//BCTCH_SIZE

设置AUTOTUNE

1
auto = tf.data.experimental.AUTOTUNE

将加载函数map到dataset

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

设置train、val数据集batch等

1
2
dataset_train = dataset_train.shuffle(BUFFER_SIZE).batch(BCTCH_SIZE)
dataset_val = dataset_val.batch(BCTCH_SIZE)

定义LinkNet模型

编写卷积模块

编写卷积模块: 1.conv 2.bn 3.relu

1
2
3
4
5
6
7
8
9
10
11
12
13
class ConvBlock(keras.layers.Layer):
    def __init__(self,units,k_size=3,stride=1,use_bias=False):
        super(ConvBlock,self).__init__()
        self.conv = keras.layers.Conv2D(units,
                                        kernel_size=k_size,
                                        strides=stride,
                                        padding='same',
                                        use_bias=use_bias)
        self.bn = keras.layers.BatchNormalization()
    def call(self,x):
        x = self.conv(x)
        x = tf.nn.relu(self.bn(x))
        return x

编写反卷积模块

编写反卷积模块: 1.conv 2.bn 3.act

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class DeconvBlock(keras.layers.Layer):
    def __init__(self,units,k_size=3,stride=2, use_bias=False):
        super(DeconvBlock,self).__init__()
        self.deconv = keras.layers.Conv2DTranspose(units,
                                        kernel_size=k_size,
                                        strides=stride,
                                        padding='same',
                                        use_bias=use_bias)
        self.bn = keras.layers.BatchNormalization()
    def call(self,x,is_act=True):
        x = self.deconv(x)
        if is_act:
            x = tf.nn.relu(self.bn(x))
        return x

编写encode block

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class EncodeBlock(keras.layers.Layer):
    def __init__(self,units):
        super(EncodeBlock,self).__init__()
        self.conv1_1 = ConvBlock(units, stride=2)
        self.conv1_2 = ConvBlock(units, stride=1)
        self.conv2_1 = ConvBlock(units, stride=1)
        self.conv2_2 = ConvBlock(units, stride=1)
        self.shortcut = ConvBlock(units, stride=2)
    def call(self, x):
        out1 = self.conv1_1(x)
        out1 = self.conv2_1(out1)
        residue = self.shortcut(x)
        out2 = self.conv2_1(out1 + residue)
        out2 = self.conv2_2(out2)
        return out2 + out1

编写Decode block

1
2
3
4
5
6
7
8
9
10
11
class DecodeBlock(keras.layers.Layer):
    def __init__(self, units1, units2):
        super(DecodeBlock, self).__init__()
        self.conv1 = ConvBlock(units1, k_size=1)
        self.deconv = DeconvBlock(units1)
        self.conv2 = ConvBlock(units2, k_size=1)
    def call(self, x):
        x = self.conv1(x)
        x = self.deconv(x)
        x = self.conv2(x)
        return x

编写Net

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
class Net(keras.Model):
    def __init__(self):
        super(Net,self).__init__()
        self.input_conv = ConvBlock(64, k_size=3, stride=2)#256X256  -> 128
        self.input_maxpool = tf.keras.layers.MaxPooling2D(padding='same')#64
        self.encode1 = EncodeBlock(64)#32
        self.encode2 = EncodeBlock(128)#16
        self.encode3 = EncodeBlock(256)#8
        self.encode4 = EncodeBlock(512)#4     4*4*512
        
        self.decode4 =DecodeBlock(512//4,256) #8*8*256
        self.decode3 =DecodeBlock(256//4,128)
        self.decode2 =DecodeBlock(128//4,64)
        self.decode1 =DecodeBlock(64//4,64)
        
        self.deconv_last1 = DeconvBlock(32)
        self.conv_last = ConvBlock(32)
        self.deconv_last2 = DeconvBlock(34,k_size=2)
    def call(self, x):
        x = self.input_conv(x)              #  (6, 128, 128, 64)
        x = self.input_maxpool(x)           #  (6, 64, 64, 64)             
        e1 = self.encode1(x)               #  (6, 32, 32, 64)
        e2 = self.encode2(e1)              #  (6, 16, 16, 128)
        e3 = self.encode3(e2)              #  (6, 8, 8, 256)
        e4 = self.encode4(e3)              #  (6, 4, 4, 512)
        
        d4 = self.decode4(e4) + e3
        d3 = self.decode3(d4) + e2
        d2 = self.decode2(d3) + e1
        d1 = self.decode1(d2)
        
        l1 = self.deconv_last1(d1)
        l2 = self.conv_last(l1)
        l3 = self.deconv_last2(l2,is_act=False)
        
        return l3

模型调用

1
model = Net()

绘制模型数据流向

1
tf.keras.utils.plot_model(model)                     # conda install pydot

设置IoU

1
2
3
4
class MeanIoU(tf.keras.metrics.MeanIoU):
    def __call__(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.argmax(y_pred, axis=-1)
        return super().__call__(y_true, y_pred, sample_weight=sample_weight)

设置optimizer

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

设置损失(SparseCategoricalCrossentropy)

1
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

设置train、test的loss、accuracy、iou

1
2
3
4
5
6
7
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
train_iou = MeanIoU(34, name='train_iou')

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

设置train数据每步

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_object(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数据每步

1
2
3
4
5
6
7
8
@tf.function
def test_step(images, labels):
    predictions = model(images)
    t_loss = loss_object(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
27
EPOCHS = 80

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 {:3d}, 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
model.save_weights('city_weight_LinkNet.h5')

进行图片预测

1
2
3
4
5
6
7
8
9
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_img =test_png
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_img.numpy()))
plt.subplot(1,2,2)
plt.imshow(np.squeeze(test_p))#squeeze 函数：从数组的形状中删除单维度条目，即把shape中为1的维度去掉