import os import cv2 import numpy as np import tensorflow as tf from tensorflow import keras import matplotlib.pyplot as plt import matplotlib as mpl from tqdm import tqdm from sklearn.model_selection import train_test_split
# a list to collect paths of 1000 images image_path = [] for root, dirs, files in os.walk('/content/png_images'): # iterate over 1000 images for file in files: # create path path = os.path.join(root,file) # add path to list image_path.append(path) len(image_path)
# a list to collect paths of 1000 masks mask_path = [] for root, dirs, files in os.walk('/content/png_masks'): #iterate over 1000 masks for file in files: # obtain the path path = os.path.join(root,file) # add path to the list mask_path.append(path) len(mask_path)
# create a list to store images images = [] # iterate over 1000 image paths for path in tqdm(image_path): # read file file = tf.io.read_file(path) # decode png file into a tensor image = tf.image.decode_png(file, channels=3, dtype=tf.uint8) # append to the list images.append(image)
# create a list to store masks masks = [] # iterate over 1000 mask paths for path in tqdm(mask_path): # read the file file = tf.io.read_file(path) # decode png file into a tensor mask = tf.image.decode_png(file, channels=1, dtype=tf.uint8) # append mask to the list masks.append(mask)
可视化数据集样本
下面的代码使用 matplotlib 使用 for 循环创建 4 到 6 范围内的图像图形。
plt.figure(figsize=(25,13))
# Iterate over the images in the range 4-6 for i in range(4,7): # Create a subplot for each image plt.subplot(4,6,i) # Get the i-th image from the list img = images[i] # Show the image with a colorbar plt.imshow(img) plt.colorbar() # Turn off the axis labels plt.axis('off')
# Display the figure plt.show()
输出：
我们再次使用 matplotlib 打印出相应的蒙版。我们定义了一个规范化器，以便蒙版具有一致性。
# Define a normalizer that can be applied while visualizing masks to have a consistency NORM = mpl.colors.Normalize(vmin=0, vmax=58)
# plot masks plt.figure(figsize=(25,13)) for i in range(4,7): plt.subplot(4,6,i) img = masks[i] plt.imshow(img, cmap='jet', norm=NORM) plt.colorbar() plt.axis('off') plt.show()
数据预处理
我们将使用两个函数开始数据预处理，以从数据集中获取图像和相应的蒙版，并将它们的大小调整为 128 x 128 像素的固定大小。
# adjust brightness of image # don't alter in mask defbrightness(img, mask): img = tf.image.adjust_brightness(img, 0.1) return img, mask
# adjust gamma of image # don't alter in mask defgamma(img, mask): img = tf.image.adjust_gamma(img, 0.1) return img, mask
# adjust hue of image # don't alter in mask defhue(img, mask): img = tf.image.adjust_hue(img, -0.1) return img, mask
defcrop(img, mask): # crop both image and mask identically img = tf.image.central_crop(img, 0.7) # resize after cropping img = tf.image.resize(img, (128,128)) mask = tf.image.central_crop(mask, 0.7) # resize afer cropping mask = tf.image.resize(mask, (128,128)) # cast to integers as they are class numbers mask = tf.cast(mask, tf.uint8) return img, mask # flip both image and mask identically defflip_hori(img, mask): img = tf.image.flip_left_right(img) mask = tf.image.flip_left_right(mask) return img, mask
# flip both image and mask identically defflip_vert(img, mask): img = tf.image.flip_up_down(img) mask = tf.image.flip_up_down(mask) return img, mask
# rotate both image and mask identically defrotate(img, mask): img = tf.image.rot90(img) mask = tf.image.rot90(mask) return img, mask
然后我们将解压缩图像和蒙版文件，应用增强函数，并将新数据连接到训练集。
# zip images and masks train = tf.data.Dataset.zip((train_X, train_y)) val = tf.data.Dataset.zip((val_X, val_y))
# perform augmentation on train data only
a = train.map(brightness) b = train.map(gamma) c = train.map(hue) d = train.map(crop) e = train.map(flip_hori) f = train.map(flip_vert) g = train.map(rotate)
#output of these layers skip_outputs = [base.get_layer(name).output for name in skip_names] #Building the downstack with the above layers. We use the pre-trained model as such, without any fine-tuning. downstack = keras.Model(inputs=base.input, outputs=skip_outputs) # freeze the downstack layers downstack.trainable = False
# define the input layer inputs = keras.layers.Input(shape=[128,128,3])
# downsample down = downstack(inputs) out = down[-1]
# prepare skip-connections skips = reversed(down[:-1]) # choose the last layer at first 4 --> 8
# upsample with skip-connections for up, skip in zip(upstack,skips): out = up(out) out = keras.layers.Concatenate()([out,skip])
# define the final transpose conv layer # image 128 by 128 with 59 classes out = keras.layers.Conv2DTranspose(59, 3, strides=2, padding='same', )(out) # complete unet model unet = keras.Model(inputs=inputs, outputs=out)
编译和训练模型
编译模型的函数，学习率为 0.001，精度为评估指标。
# compiling the model defCompile_Model(): unet.compile(loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=keras.optimizers.RMSprop(learning_rate=0.001), metrics=['accuracy']) Compile_Model()
# select a validation data batch img, mask = next(iter(val)) # make prediction pred = unet.predict(img) plt.figure(figsize=(20,28))
k = 0 for i in pred: # plot the predicted mask plt.subplot(4,3,1+k*3) i = tf.argmax(i, axis=-1) plt.imshow(i,cmap='jet', norm=NORM) plt.axis('off') plt.title('Prediction')
# plot the actual image plt.subplot(4,3,3+k*3) plt.imshow(img[k]) plt.axis('off') plt.title('Actual Image') k += 1 if k == 4: break plt.suptitle('Predition After 20 Epochs (No Fine-tuning)', color='red', size=20) plt.show()
从第 21 个 epoch 到第 40 个 epoch 训练模型和微调。
downstack.trainable = True # compile again Compile_Model() # train from epoch 20 to 40 hist_2 = unet.fit(train, validation_data=val, steps_per_epoch=STEPS_PER_EPOCH, validation_steps=VALIDATION_STEPS, epochs=40, initial_epoch = 20, verbose = 2 )
使用模型进行蒙版预测。
# select a validation data batch img, mask = next(iter(val)) # make prediction pred = unet.predict(img) plt.figure(figsize=(20,30))
k = 0 for i in pred: # plot the predicted mask plt.subplot(4,3,1+k*3) i = tf.argmax(i, axis=-1) plt.imshow(i,cmap='jet', norm=NORM) plt.axis('off') plt.title('Prediction')
# plot the actual image plt.subplot(4,3,3+k*3) plt.imshow(img[k]) plt.axis('off') plt.title('Actual Image') k += 1 if k == 4: break plt.suptitle('Predition After 40 Epochs (By Fine-tuning from 21th Epoch)', color='red', size=20) plt.show()
