from tensorflow.keras.layers import (Input, Conv2D, MaxPooling2D,
UpSampling2D,Concatenate)
from tensorflow.keras.models import Model
img_rows = 512
img_cols = 512
def UNET():
inputs = Input(shape=(img_rows,img_cols,1),name='initial_input')
conv1 = Conv2D(32, 3, activation='relu', padding='same',name='NO1_conv_1')(inputs)
conv1 = Conv2D(32, 3, activation='relu', padding='same',name='NO1_conv_2')(conv1)
pool1 = MaxPooling2D(name='NO1_maxpooling')(conv1)
conv2 = Conv2D(64, 3, activation='relu', padding='same',name='NO2_conv_1')(pool1)
conv2 = Conv2D(64, 3, activation='relu', padding='same',name='NO2_conv_2')(conv2)
pool2 = MaxPooling2D(name='NO2_maxpooling')(conv2)
conv3 = Conv2D(128, 3, activation='relu', padding='same',name='NO3_conv_1')(pool2)
conv3 = Conv2D(128, 3, activation='relu', padding='same',name='NO3_conv_2')(conv3)
pool3 = MaxPooling2D(name='NO3_maxpooling')(conv3)
conv4 = Conv2D(256, 3, activation='relu', padding='same',name='NO4_conv_1')(pool3)
conv4 = Conv2D(256, 3, activation='relu', padding='same',name='NO4_conv_2')(conv4)
pool4 = MaxPooling2D(name='NO4_maxpooling')(conv4)
conv5 = Conv2D(512, 3, activation='relu', padding='same',name='NO5_conv_1')(pool4)
conv5 = Conv2D(512, 3, activation='relu', padding='same',name='NO5_conv_2')(conv5)
conv5_up = UpSampling2D(name='conv5_up')(conv5)
up6 = Concatenate(axis=3,name='up6')([conv5_up,conv4])
conv6 = Conv2D(256, 3, activation='relu', padding='same',name='NO6_conv_1')(up6)
conv6 = Conv2D(256, 3, activation='relu', padding='same',name='NO6_conv_2')(conv6)
conv6_up = UpSampling2D(name='conv6_up')(conv6)
up7 = Concatenate(axis=3,name='up7')([conv6_up,conv3])
conv7 = Conv2D(128,3, activation='relu', padding='same',name='NO7_conv_1')(up7)
conv7 = Conv2D(128,3, activation='relu', padding='same',name='NO7_conv_2')(conv7)
conv7_up = UpSampling2D(name='conv7_up')(conv7)
up8 = Concatenate(axis=3,name='up8')([conv7_up,conv2])
conv8 = Conv2D(64,3, activation='relu', padding='same',name='NO8_conv_1')(up8)
conv8 = Conv2D(64,3, activation='relu', padding='same',name='NO8_conv_2')(conv8)
conv8_up = UpSampling2D(name='conv8_up')(conv8)
up9 = Concatenate(axis=3,name='up9')([conv8_up,conv1])
conv9 = Conv2D(32, 3, activation='relu', padding='same',name='NO9_conv_1')(up9)
conv9 = Conv2D(32, 3, activation='relu', padding='same',name='NO9_conv_2')(conv9)
conv10 = Conv2D(1, 1, activation='sigmoid',name='cls_layer')(conv9) # num_classes=1, 1*1卷积,相当于dense层
model = Model(inputs=inputs, outputs=conv10,name='my_unet_tf2')
model.summary()
if __name__ == '__main__':
UNET()
Model: "my_unet_tf2"
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
initial_input (InputLayer) [(None, 512, 512, 1 0 []
)]
NO1_conv_1 (Conv2D) (None, 512, 512, 32 320 ['initial_input[0][0]']
)
NO1_conv_2 (Conv2D) (None, 512, 512, 32 9248 ['NO1_conv_1[0][0]']
)
NO1_maxpooling (MaxPooling2D) (None, 256, 256, 32 0 ['NO1_conv_2[0][0]']
)
NO2_conv_1 (Conv2D) (None, 256, 256, 64 18496 ['NO1_maxpooling[0][0]']
)
NO2_conv_2 (Conv2D) (None, 256, 256, 64 36928 ['NO2_conv_1[0][0]']
)
NO2_maxpooling (MaxPooling2D) (None, 128, 128, 64 0 ['NO2_conv_2[0][0]']
)
NO3_conv_1 (Conv2D) (None, 128, 128, 12 73856 ['NO2_maxpooling[0][0]']
8)
NO3_conv_2 (Conv2D) (None, 128, 128, 12 147584 ['NO3_conv_1[0][0]']
8)
NO3_maxpooling (MaxPooling2D) (None, 64, 64, 128) 0 ['NO3_conv_2[0][0]']
NO4_conv_1 (Conv2D) (None, 64, 64, 256) 295168 ['NO3_maxpooling[0][0]']
NO4_conv_2 (Conv2D) (None, 64, 64, 256) 590080 ['NO4_conv_1[0][0]']
NO4_maxpooling (MaxPooling2D) (None, 32, 32, 256) 0 ['NO4_conv_2[0][0]']
NO5_conv_1 (Conv2D) (None, 32, 32, 512) 1180160 ['NO4_maxpooling[0][0]']
NO5_conv_2 (Conv2D) (None, 32, 32, 512) 2359808 ['NO5_conv_1[0][0]']
conv5_up (UpSampling2D) (None, 64, 64, 512) 0 ['NO5_conv_2[0][0]']
up6 (Concatenate) (None, 64, 64, 768) 0 ['conv5_up[0][0]',
'NO4_conv_2[0][0]']
NO6_conv_1 (Conv2D) (None, 64, 64, 256) 1769728 ['up6[0][0]']
NO6_conv_2 (Conv2D) (None, 64, 64, 256) 590080 ['NO6_conv_1[0][0]']
conv6_up (UpSampling2D) (None, 128, 128, 25 0 ['NO6_conv_2[0][0]']
6)
up7 (Concatenate) (None, 128, 128, 38 0 ['conv6_up[0][0]',
4) 'NO3_conv_2[0][0]']
NO7_conv_1 (Conv2D) (None, 128, 128, 12 442496 ['up7[0][0]']
8)
NO7_conv_2 (Conv2D) (None, 128, 128, 12 147584 ['NO7_conv_1[0][0]']
8)
conv7_up (UpSampling2D) (None, 256, 256, 12 0 ['NO7_conv_2[0][0]']
8)
up8 (Concatenate) (None, 256, 256, 19 0 ['conv7_up[0][0]',
2) 'NO2_conv_2[0][0]']
NO8_conv_1 (Conv2D) (None, 256, 256, 64 110656 ['up8[0][0]']
)
NO8_conv_2 (Conv2D) (None, 256, 256, 64 36928 ['NO8_conv_1[0][0]']
)
conv8_up (UpSampling2D) (None, 512, 512, 64 0 ['NO8_conv_2[0][0]']
)
up9 (Concatenate) (None, 512, 512, 96 0 ['conv8_up[0][0]',
) 'NO1_conv_2[0][0]']
NO9_conv_1 (Conv2D) (None, 512, 512, 32 27680 ['up9[0][0]']
)
NO9_conv_2 (Conv2D) (None, 512, 512, 32 9248 ['NO9_conv_1[0][0]']
)
cls_layer (Conv2D) (None, 512, 512, 1) 33 ['NO9_conv_2[0][0]']
==================================================================================================
Total params: 7,846,081
Trainable params: 7,846,081
Non-trainable params: 0
__________________________________________________________________________________________________
Process finished with exit code 0