代码
导包
import tensorflow as tf
import glob
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from IPython.display import display
import os.path
常量
# 图片大小
IMG_WIDTH=224
AUTOTUNE=tf.data.experimental.AUTOTUNE
# vgg16权重文件
vgg16_h5='/kaggle/input/vgg16-weights-tf-dim-ordering-tf-kernels-notop/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'
获取所有jpg图片路径
# 有3个mat文件
img_paths=glob.glob("/kaggle/input/head-location/images/images/*.jpg")
# png_paths=glob.glob("/kaggle/input/head-location/annotations/annotations/trimaps/*")
print(img_paths[:3])
# print(png_paths[:3])
查看所有图片数量
count_img=len(img_paths)
# count_png=len(png_paths)
print(count_img)
# print(count_png)
获取所有png图片路径
jpg图片和png图片可能不是一一对应的，所以需要根据jpg图片名字，获取png图片路径
png_path="/kaggle/input/head-location/annotations/annotations/trimaps/"
png_paths=[]
def get_png_paths():
for item in img_paths:
# print(item)
name=os.path.basename(item).split(".")[0]+".png"
# print(name)
path=os.path.join(png_path,name)
png_paths.append(path)
# print(path)
get_png_paths()
print(png_paths[0])
print(len(png_paths))
测试：显示jpg和png图片
display(Image.open(img_paths[0]))
display(Image.open(png_paths[0]))
注意： png图片不能正常显示
查看png图片内部数据（关键）
from collections import Counter
def show_png_data():
arr=np.array(Image.open(png_paths[0]))
print(type(arr))
print(arr)
print("最小数：",np.min(arr))
print("最大数：",np.max(arr))
print(Counter(arr.flatten()))
show_png_data()
执行结果：
<class 'numpy.ndarray'>
[[2 2 2 ... 2 2 2]
[2 2 2 ... 2 2 2]
[2 2 2 ... 2 2 2]
...
[2 2 2 ... 2 2 2]
[2 2 2 ... 2 2 2]
[2 2 2 ... 2 2 2]]
最小数： 1
最大数： 3
Counter({2: 132176, 1: 37847, 3: 17477})
可知png图片的数据只有1，2，3
用matplotlib图片
matplotlib显示特殊图片时，会自行调整，显示成人眼可识别的图片
plt.imshow(Image.open(img_paths[0]))
plt.show()
plt.imshow(Image.open(png_paths[0]))
plt.show()
生成Dataset
ds=tf.data.Dataset.from_tensor_slices((img_paths,png_paths))
测试
for jpg,png in ds.take(2):
print(jpg)
print(png)
print('--')
实现解析图片函数（关键）
def parse_jpg(path):
img=tf.io.read_file(path)
img=tf.image.decode_jpeg(img,channels=3)
img=tf.image.resize(img,(IMG_WIDTH,IMG_WIDTH))
img=img/255
return img
下面代码处理png图片时，要-1（这是laben）
def parse_png(path):
img=tf.io.read_file(path)
img=tf.image.decode_png(img,channels=1)
img=tf.image.resize(img,(IMG_WIDTH,IMG_WIDTH))
img=img-1
return img
def parse_img(jpg_path,png_path):
jpg=parse_jpg(jpg_path)
png=parse_png(png_path)
return jpg,png
执行函数
ds2=ds.map(parse_img,num_parallel_calls=AUTOTUNE).shuffle(count_img)
分割训练集和测试集
train_count=int(count_img*0.7)
train_ds=ds2.take(train_count)
test_ds=ds2.skip(train_count)
测试图片
for jpg,png in train_ds.take(2):
# print(jpg.numpy())
# arr=png.numpy()
# print(collections.Counter(arr.flatten()))
plt.imshow(jpg.numpy())
plt.show()
png2=np.squeeze(png.numpy())
plt.imshow(png2)
plt.show()
准备训练集和测试集
train_ds2=train_ds.shuffle(train_count).batch(8).prefetch(AUTOTUNE)
test_ds2=test_ds.batch(8).prefetch(AUTOTUNE)
train_ds
加载vgg16模型
使用vgg16模型提取图片的特征，所以设置其不可训练
app_vgg16=tf.keras.applications.VGG16(include_top=False,weights=vgg16_h5,input_shape=(IMG_WIDTH,IMG_WIDTH,3))
app_vgg16.trainable=False
app_vgg16.summary()
构建FCN模型（关键）
def build_fcn8_vgg16():
pool3=app_vgg16.get_layer("block3_pool").output
print(pool3.shape)
pool4=app_vgg16.get_layer("block4_pool").output
print(pool4.shape)
# 卷积层
print(app_vgg16.name,":",app_vgg16.output.shape)
# 上采样，放大2倍
o=tf.keras.layers.Conv2DTranspose(512,2,2,padding="same",activation="relu",name="myblock8_conv_convtran")(app_vgg16.output) #padding="same",
print(o.name,":",o.shape)
# 对pool4进行卷积运算
o2=tf.keras.layers.Conv2D(512,1,1,activation="relu",padding="same",name="pool4_conv")(pool4)
print(o2.shape)
o=tf.add(o,o2,name="myblock8_conv_convtran_add_pool4_conv")
# 上采样，放大2倍
o = tf.keras.layers.Conv2DTranspose(256, 2, 2, padding="same", activation="relu",name="sum_convtran")(o)
# 对pool3进行卷积运算
o2=tf.keras.layers.Conv2D(256,1,1,activation="relu",padding="same",name="pool3_conv")(pool3)
print(o2.shape)
o=tf.add(o,o2,name="sum_convtran_add_pool3_conv")
o=tf.keras.layers.Conv2D(64,3,1,padding="same",activation="relu")(o)
print(o.shape)
# 对最终结果放大8倍
o = tf.keras.layers.Conv2DTranspose(3, 8, 8, padding="same", activation="softmax",name="result")(o)
model=tf.keras.Model(inputs=app_vgg16.input,outputs=[o],name="fcn8s_vgg16")
print("|"*50)
model.summary()
return model
model=build_fcn8_vgg16()
执行结果：
(None, 28, 28, 256)
(None, 14, 14, 512)
vgg16 : (None, 7, 7, 512)
myblock8_conv_convtran/Identity:0 : (None, 14, 14, 512)
(None, 14, 14, 512)
(None, 28, 28, 256)
(None, 28, 28, 64)
||||||||||||||||||||||||||||||||||||||||||||||||||
Model: "fcn8s_vgg16"
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input_1 (InputLayer) [(None, 224, 224, 3) 0
__________________________________________________________________________________________________
block1_conv1 (Conv2D) (None, 224, 224, 64) 1792 input_1[0][0]
__________________________________________________________________________________________________
block1_conv2 (Conv2D) (None, 224, 224, 64) 36928 block1_conv1[0][0]
__________________________________________________________________________________________________
block1_pool (MaxPooling2D) (None, 112, 112, 64) 0 block1_conv2[0][0]
__________________________________________________________________________________________________
block2_conv1 (Conv2D) (None, 112, 112, 128 73856 block1_pool[0][0]
__________________________________________________________________________________________________
block2_conv2 (Conv2D) (None, 112, 112, 128 147584 block2_conv1[0][0]
__________________________________________________________________________________________________
block2_pool (MaxPooling2D) (None, 56, 56, 128) 0 block2_conv2[0][0]
__________________________________________________________________________________________________
block3_conv1 (Conv2D) (None, 56, 56, 256) 295168 block2_pool[0][0]
__________________________________________________________________________________________________
block3_conv2 (Conv2D) (None, 56, 56, 256) 590080 block3_conv1[0][0]
__________________________________________________________________________________________________
block3_conv3 (Conv2D) (None, 56, 56, 256) 590080 block3_conv2[0][0]
__________________________________________________________________________________________________
block3_pool (MaxPooling2D) (None, 28, 28, 256) 0 block3_conv3[0][0]
__________________________________________________________________________________________________
block4_conv1 (Conv2D) (None, 28, 28, 512) 1180160 block3_pool[0][0]
__________________________________________________________________________________________________
block4_conv2 (Conv2D) (None, 28, 28, 512) 2359808 block4_conv1[0][0]
__________________________________________________________________________________________________
block4_conv3 (Conv2D) (None, 28, 28, 512) 2359808 block4_conv2[0][0]
__________________________________________________________________________________________________
block4_pool (MaxPooling2D) (None, 14, 14, 512) 0 block4_conv3[0][0]
__________________________________________________________________________________________________
block5_conv1 (Conv2D) (None, 14, 14, 512) 2359808 block4_pool[0][0]
__________________________________________________________________________________________________
block5_conv2 (Conv2D) (None, 14, 14, 512) 2359808 block5_conv1[0][0]
__________________________________________________________________________________________________
block5_conv3 (Conv2D) (None, 14, 14, 512) 2359808 block5_conv2[0][0]
__________________________________________________________________________________________________
block5_pool (MaxPooling2D) (None, 7, 7, 512) 0 block5_conv3[0][0]
__________________________________________________________________________________________________
myblock8_conv_convtran (Conv2DT (None, 14, 14, 512) 1049088 block5_pool[0][0]
__________________________________________________________________________________________________
pool4_conv (Conv2D) (None, 14, 14, 512) 262656 block4_pool[0][0]
__________________________________________________________________________________________________
tf_op_layer_myblock8_conv_convt [(None, 14, 14, 512) 0 myblock8_conv_convtran[0][0]
pool4_conv[0][0]
__________________________________________________________________________________________________
sum_convtran (Conv2DTranspose) (None, 28, 28, 256) 524544 tf_op_layer_myblock8_conv_convtra
__________________________________________________________________________________________________
pool3_conv (Conv2D) (None, 28, 28, 256) 65792 block3_pool[0][0]
__________________________________________________________________________________________________
tf_op_layer_sum_convtran_add_po [(None, 28, 28, 256) 0 sum_convtran[0][0]
pool3_conv[0][0]
__________________________________________________________________________________________________
conv2d (Conv2D) (None, 28, 28, 64) 147520 tf_op_layer_sum_convtran_add_pool
__________________________________________________________________________________________________
result (Conv2DTranspose) (None, 224, 224, 3) 12291 conv2d[0][0]
==================================================================================================
Total params: 16,776,579
Trainable params: 2,061,891
Non-trainable params: 14,714,688
__________________________________________________________________________________________________
可视化模型
tf.keras.utils.plot_model(model2,to_file="model.png",show_shapes=True)
编译 和 训练
def exe_model():
# model.compile(optimizer="adam",loss="sparse_categorical_crossentropy",metrics=["acc"])
model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.001),loss="sparse_categorical_crossentropy",metrics=["acc"])
history=model.fit(train_ds2,epochs=20,validation_data=test_ds2)
return history
history=exe_model()
查看损失值 和 评价值
plt.plot(history.epoch,history.history.get("loss"),label="loss")
plt.plot(history.epoch,history.history.get("val_loss"),label="val_loss")
plt.legend()
plt.show()
plt.plot(history.epoch,history.history.get("acc"),label="acc")
plt.plot(history.epoch,history.history.get("val_acc"),label="val_acc")
plt.legend()
plt.show()
保存模型
model.save("model.h5")