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

打印Tensorflow版本

1
print('Tensorflow version:{}'.format(tf.__version__))

打印是否支持GPU加速

1
print(tf.test.is_gpu_available())

获取图片路径

1
images = glob.glob('./images/*.jpg')

获取标注图片的路径

1
anno = glob.glob('./annotations/trimaps/*.png')

图片和标注图片做乱序

1
2
3
4
np.random.seed(2023)
index = np.random.permutation(len(images))
images = np.array(images)[index]
anno = np.array(anno)[index]

img = tf.squeeze(img)#自动把通道数为1的压缩成2二维列表

图片和标注图片Dataset

1
dataset = tf.data.Dataset.from_tensor_slices((images,anno))

划分数据集

1
2
3
4
test_count = int(len(images)*0.2)
train_count = len(images) - test_count
data_train = dataset.skip(test_count)
data_test = dataset.take(test_count)

图片和标注图片的读取函数

1
2
3
4
def read_jpg(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img,channels=3)
    return img
1
2
3
4
def read_png(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img,channels=1)
    return img
1
2
3
4
5
def normal_img(input_images,input_anno):
    input_images = tf.cast(input_images,tf.float32)
    input_images = input_images/127.5 - 1 #[-1.1]
    input_anno = input_anno - 1
    return input_images,input_anno
1
2
3
4
5
6
def load_images(input_images_path,input_anno_path):
    input_image = read_jpg(input_images_path)
    input_anno = read_png(input_anno_path)
    input_image = tf.image.resize(input_image,(224,224))
    input_anno = tf.image.resize(input_anno,(224,224))
    return normal_img(input_image,input_anno)

Dataset加载图片读取函数

1
2
data_train = data_train.map(load_images,num_parallel_calls=tf.data.experimental.AUTOTUNE)
data_test = data_test.map(load_images,num_parallel_calls=tf.data.experimental.AUTOTUNE)

设置BATCH_SIZE

1
BATCH_SIZE = 32

对Dataset设置repeat、shuffle、batch

1
2
data_train = data_train.repeat().shuffle(100).batch(BATCH_SIZE)#之前路径乱序了，这里可以小点
data_test = data_test.batch(BATCH_SIZE)

查看Dataset中图片与标注图片

1
2
3
4
5
for img,anno in data_train.take(1):
    plt.subplot(1,2,1)#绘制子图  表示分成1行2列 目前第一个
    plt.imshow(tf.keras.preprocessing.image.array_to_img(img[0]))#将tensor转换成arrary
    plt.subplot(1,2,2)#绘制子图  表示分成1行2列 目前第二个
    plt.imshow(tf.keras.preprocessing.image.array_to_img(anno[0]))#将tensor转换成arrury

加载VGG16预训练模型

1
2
3
4
conv_base = tf.keras.applications.VGG16(weights='imagenet',
                                       input_shape=(224,224,3),
                                       include_top=False
)

image-20230201082534222

获取中间层的输出值

1
2
3
4
5
6
layer_names = [
    'block5_conv3',
    'block4_conv3',
    'block3_conv3',
    'block5_pool'#原始最后一层
]
1
layers_output = [conv_base.get_layer(layer_name).output for layer_name in layer_names]

创建含有多个输出的子模型

1
2
3
multi_out_model = tf.keras.models.Model(inputs=conv_base.input,
                                 outputs=layers_output
)#创建含有多个输出的子模型

设置模型为不可训练（保留训练权重）

1
multi_out_model.trainable = False

input模型得到输出

1
2
inputs = tf.keras.layers.Input(shape=(224,224,3))
out_block5_conv3,out_block4_conv3,out_block3_conv3,out = multi_out_model(inputs)

out_block5_conv3.shape,out_block4_conv3.shape,out_block3_conv3.shape,out.shape (TensorShape([None, 14, 14, 512]), TensorShape([None, 28, 28, 512]), TensorShape([None, 56, 56, 256]), TensorShape([None, 7, 7, 512]))

上采样 Upsampling 反卷积（转置卷积）

1
2
3
4
5
6
7
x1 = tf.keras.layers.Conv2DTranspose(filters=512,
                                     kernel_size=3,
                                     strides=2,
                                     padding='same',
                                    activation='relu')(out)#反卷积（转置卷积） valid尽量不填充  same填充为原来2倍   进行上采样
# 加层卷积提取特征
x1 = tf.keras.layers.Conv2D(512,3,padding='same',activation='relu')(x1)

TensorShape([None, 14, 14, 512])

1
2
3
4
5
6
7
8
x2 = tf.add(x1,out_block5_conv3)#tensor元素相加
x2 = tf.keras.layers.Conv2DTranspose(filters=512,
                                     kernel_size=3,
                                     strides=2,
                                     padding='same',
                                    activation='relu')(x2)#反卷积（转置卷积） valid尽量不填充  same填充为原来2倍   进行上采样
# 加层卷积提取特征
x2 = tf.keras.layers.Conv2D(512,3,padding='same',activation='relu')(x2)

TensorShape([None, 28, 28, 512])

1
2
3
4
5
6
7
8
x3 = tf.add(x2,out_block4_conv3)#tensor元素相加
x3 = tf.keras.layers.Conv2DTranspose(filters=256,
                                     kernel_size=3,
                                     strides=2,
                                     padding='same',
                                    activation='relu')(x3)#反卷积（转置卷积） valid尽量不填充  same填充为原来2倍   进行上采样
# 加层卷积提取特征
x3 = tf.keras.layers.Conv2D(256,3,padding='same',activation='relu')(x3)

TensorShape([None, 56, 56, 256])

1
2
3
4
5
6
7
8
x4 = tf.add(x3,out_block3_conv3)#tensor元素相加
x5 = tf.keras.layers.Conv2DTranspose(filters=128,
                                     kernel_size=3,
                                     strides=2,
                                     padding='same',
                                    activation='relu')(x4)#反卷积（转置卷积） valid尽量不填充  same填充为原来2倍   进行上采样
# 加层卷积提取特征
x5 = tf.keras.layers.Conv2D(128,3,padding='same',activation='relu')(x5)

TensorShape([None, 112, 112, 128])

1
2
3
4
5
prediction = tf.keras.layers.Conv2DTranspose(filters=3,#3种分类
                                     kernel_size=3,
                                     strides=2,
                                     padding='same',
                                    activation='softmax')(x5)#反卷积（转置卷积） valid尽量不填充  same填充为原来2倍   进行上采样)

TensorShape([None, 224, 224, 3])

设置模型的inputs与outputs

1
2
3
4
model = tf.keras.models.Model(
    inputs=inputs,	#inputs = tf.keras.layers.Input(shape=(224,224,3))
    outputs=prediction
)

配置模型

1
2
3
4
5
model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss=tf.keras.losses.sparse_categorical_crossentropy,#数字编码
    metrics=['acc']#tf.keras.metrics.Mean('accuracy')
)

模型训练

1
2
3
4
5
6
history = model.fit(data_train,
         epochs=20,
         steps_per_epoch=train_count//BATCH_SIZE,
         validation_data=data_test,
         validation_steps=test_count//BATCH_SIZE
)

绘制模型训练过程中train_loss与test_loss变化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
loss = history.history['loss']
val_loss = history.history['val_loss']

epochs =range(20)

plt.figure()
plt.plot(epochs,loss,'r',label='Training loss')
plt.plot(epochs,val_loss,'bo',label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0,1])
plt.legend()
plt.show()

绘制原图片、标注图片、预测图片

1
2
3
4
5
6
7
8
9
10
11
12
13
14
num = 3
for image,mask in data_test.take(1):
    pred_mask = model.predict(image)
    pred_mask = tf.argmax(pred_mask,axis=-1)
    pred_mask = pred_mask[...,tf.newaxis]#维度扩展[224,244,1]   ...取前面所有维度
    
    plt.figure(figsize=(10,10))
    for i in range(num):
        plt.subplot(num,3,i*num+1)
        plt.imshow(tf.keras.preprocessing.image.array_to_img(image[i]))
        plt.subplot(num,3,i*num+2)
        plt.imshow(tf.keras.preprocessing.image.array_to_img(mask[i]))
        plt.subplot(num,3,i*num+3)
        plt.imshow(tf.keras.preprocessing.image.array_to_img(pred_mask[i]))