来源:OSC开源社区
模型摘要
# 搭建补全网络
def generator(x):
# conv1
conv1 = fluid.layers.conv2d(input=x,num_filters= 64,filter_size= 5,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv1',data_format= 'NHWC')
conv1 = fluid.layers.batch_norm(conv1, momentum= 0.99, epsilon= 0.001)
conv1 = fluid.layers.relu(conv1, name= None)
# conv2
conv2 = fluid.layers.conv2d(input=conv1,num_filters= 128,filter_size= 3,dilation= 1,stride= 2,padding= 'SAME',name= 'generator_conv2',data_format= 'NHWC')
conv2 = fluid.layers.batch_norm(conv2, momentum= 0.99, epsilon= 0.001)
conv2 = fluid.layers.relu(conv2, name= None)
# conv3
conv3 = fluid.layers.conv2d(input=conv2,num_filters= 128,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv3',data_format= 'NHWC')
conv3 = fluid.layers.batch_norm(conv3, momentum= 0.99, epsilon= 0.001)
conv3 = fluid.layers.relu(conv3, name= None)
# conv4
conv4 = fluid.layers.conv2d(input=conv3,num_filters= 256,filter_size= 3,dilation= 1,stride= 2,padding= 'SAME',name= 'generator_conv4',data_format= 'NHWC')
conv4 = fluid.layers.batch_norm(conv4, momentum= 0.99, epsilon= 0.001)
conv4 = fluid.layers.relu(conv4, name= None)
# conv5
conv5 = fluid.layers.conv2d(input=conv4,num_filters= 256,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv5',data_format= 'NHWC')
conv5 = fluid.layers.batch_norm(conv5, momentum= 0.99, epsilon= 0.001)
conv5 = fluid.layers.relu(conv5, name= None)
# conv6
conv6 = fluid.layers.conv2d(input=conv5,num_filters= 256,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv6',data_format= 'NHWC')
conv6 = fluid.layers.batch_norm(conv6, momentum= 0.99, epsilon= 0.001)
conv6 = fluid.layers.relu(conv6, name= None)
# 空洞卷积
# dilated1
dilated1 = fluid.layers.conv2d(input=conv6,num_filters= 256,filter_size= 3,dilation= 2,padding= 'SAME',name= 'generator_dilated1',data_format= 'NHWC')
dilated1 = fluid.layers.batch_norm(dilated1, momentum= 0.99, epsilon= 0.001)
dilated1 = fluid.layers.relu(dilated1, name= None)
# dilated2
dilated2 = fluid.layers.conv2d(input=dilated1,num_filters= 256,filter_size= 3,dilation= 4,padding= 'SAME',name= 'generator_dilated2',data_format= 'NHWC') #stride=1
dilated2 = fluid.layers.batch_norm(dilated2, momentum= 0.99, epsilon= 0.001)
dilated2 = fluid.layers.relu(dilated2, name= None)
# dilated3
dilated3 = fluid.layers.conv2d(input=dilated2,num_filters= 256,filter_size= 3,dilation= 8,padding= 'SAME',name= 'generator_dilated3',data_format= 'NHWC')
dilated3 = fluid.layers.batch_norm(dilated3, momentum= 0.99, epsilon= 0.001)
dilated3 = fluid.layers.relu(dilated3, name= None)
# dilated4
dilated4 = fluid.layers.conv2d(input=dilated3,num_filters= 256,filter_size= 3,dilation= 16,padding= 'SAME',name= 'generator_dilated4',data_format= 'NHWC')
dilated4 = fluid.layers.batch_norm(dilated4, momentum= 0.99, epsilon= 0.001)
dilated4 = fluid.layers.relu(dilated4, name= None)
# conv7
conv7 = fluid.layers.conv2d(input=dilated4,num_filters= 256,filter_size= 3,dilation= 1,name= 'generator_conv7',data_format= 'NHWC')
conv7 = fluid.layers.batch_norm(conv7, momentum= 0.99, epsilon= 0.001)
conv7 = fluid.layers.relu(conv7, name= None)
# conv8
conv8 = fluid.layers.conv2d(input=conv7,num_filters= 256,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv8',data_format= 'NHWC')
conv8 = fluid.layers.batch_norm(conv8, momentum= 0.99, epsilon= 0.001)
conv8 = fluid.layers.relu(conv8, name= None)
# deconv1
deconv1 = fluid.layers.conv2d_transpose(input=conv8, num_filters= 128, output_size=[ 64, 64],stride = 2,name= 'generator_deconv1',data_format= 'NHWC')
deconv1 = fluid.layers.batch_norm(deconv1, momentum= 0.99, epsilon= 0.001)
deconv1 = fluid.layers.relu(deconv1, name= None)
# conv9
conv9 = fluid.layers.conv2d(input=deconv1,num_filters= 128,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv9',data_format= 'NHWC')
conv9 = fluid.layers.batch_norm(conv9, momentum= 0.99, epsilon= 0.001)
conv9 = fluid.layers.relu(conv9, name= None)
# deconv2
deconv2 = fluid.layers.conv2d_transpose(input=conv9, num_filters= 64, output_size=[ 128, 128],stride = 2,name= 'generator_deconv2',data_format= 'NHWC')
deconv2 = fluid.layers.batch_norm(deconv2, momentum= 0.99, epsilon= 0.001)
deconv2 = fluid.layers.relu(deconv2, name= None)
# conv10
conv10 = fluid.layers.conv2d(input=deconv2,num_filters= 32,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv10',data_format= 'NHWC')
conv10 = fluid.layers.batch_norm(conv10, momentum= 0.99, epsilon= 0.001)
conv10 = fluid.layers.relu(conv10, name= None)
# conv11
x = fluid.layers.conv2d(input=conv10,num_filters= 3,filter_size= 3,dilation= 1,stride= 1,padding= 'SAME',name= 'generator_conv11',data_format= 'NHWC')
x = fluid.layers.tanh(x)
return x
# 搭建内容鉴别器
def discriminator(global_x, local_x):
def global_discriminator(x):
# conv1
conv1 = fluid.layers.conv2d(input=x,num_filters= 64,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_global_conv1',data_format= 'NHWC')
conv1 = fluid.layers.batch_norm(conv1, momentum= 0.99, epsilon= 0.001)
conv1 = fluid.layers.relu(conv1, name= None)
# conv2
conv2 = fluid.layers.conv2d(input=conv1,num_filters= 128,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_global_conv2',data_format= 'NHWC')
conv2 = fluid.layers.batch_norm(conv2, momentum= 0.99, epsilon= 0.001)
conv2 = fluid.layers.relu(conv2, name= None)
# conv3
conv3 = fluid.layers.conv2d(input=conv2,num_filters= 256,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_global_conv3',data_format= 'NHWC')
conv3 = fluid.layers.batch_norm(conv3, momentum= 0.99, epsilon= 0.001)
conv3 = fluid.layers.relu(conv3, name= None)
# conv4
conv4 = fluid.layers.conv2d(input=conv3,num_filters= 512,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_global_conv4',data_format= 'NHWC')
conv4 = fluid.layers.batch_norm(conv4, momentum= 0.99, epsilon= 0.001)
conv4 = fluid.layers.relu(conv4, name= None)
# conv5
conv5 = fluid.layers.conv2d(input=conv4,num_filters= 512,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_global_conv5',data_format= 'NHWC')
conv5 = fluid.layers.batch_norm(conv5, momentum= 0.99, epsilon= 0.001)
conv5 = fluid.layers.relu(conv5, name= None)
# conv6
conv6 = fluid.layers.conv2d(input=conv5,num_filters= 512,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_global_conv6',data_format= 'NHWC')
conv6 = fluid.layers.batch_norm(conv6, momentum= 0.99, epsilon= 0.001)
conv6 = fluid.layers.relu(conv6, name= None)
# fc
x = fluid.layers.fc(input=conv6, size= 1024,name= 'discriminator_global_fc1')
return x
def local_discriminator(x):
# conv1
conv1 = fluid.layers.conv2d(input=x,num_filters= 64,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_lobal_conv1',data_format= 'NHWC')
conv1 = fluid.layers.batch_norm(conv1, momentum= 0.99, epsilon= 0.001)
conv1 = fluid.layers.relu(conv1, name= None)
# conv2
conv2 = fluid.layers.conv2d(input=conv1,num_filters= 128,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_lobal_conv2',data_format= 'NHWC')
conv2 = fluid.layers.batch_norm(conv2, momentum= 0.99, epsilon= 0.001)
conv2 = fluid.layers.relu(conv2, name= None)
# conv3
conv3 = fluid.layers.conv2d(input=conv2,num_filters= 256,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_lobal_conv3',data_format= 'NHWC')
conv3 = fluid.layers.batch_norm(conv3, momentum= 0.99, epsilon= 0.001)
conv3 = fluid.layers.relu(conv3, name= None)
# conv4
conv4 = fluid.layers.conv2d(input=conv3,num_filters= 512,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_lobal_conv4',data_format= 'NHWC')
conv4 = fluid.layers.batch_norm(conv4, momentum= 0.99, epsilon= 0.001)
conv4 = fluid.layers.relu(conv4, name= None)
# conv5
conv5 = fluid.layers.conv2d(input=conv4,num_filters= 512,filter_size= 5,dilation= 1,stride= 2,padding= 'SAME',name= 'discriminator_lobal_conv5',data_format= 'NHWC')
conv5 = fluid.layers.batch_norm(conv5, momentum= 0.99, epsilon= 0.001)
conv5 = fluid.layers.relu(conv5, name= None)
# fc
x = fluid.layers.fc(input=conv5, size= 1024,name= 'discriminator_lobal_fc1')
return x
global_output = global_discriminator(global_x)
local_output = local_discriminator(local_x)
print( 'global_output',global_output.shape)
print( 'local_output',local_output.shape)
output = fluid.layers.concat([global_output, local_output], axis= 1)
output = fluid.layers.fc(output, size= 1,name= 'discriminator_concatenation_fc1')
return output
# 生成器优先迭代次数
NUM _TRAIN_TIMES _OF_DG = 100
# 总迭代轮次
epoch = 200
step _num = int(len(x_train) / BATCH_SIZE)
np.random.shuffle(x_train)
for pass_id in range(epoch):
# 训练生成器
if pass _id <= NUM_TRAIN _TIMES_OF_DG:
g _loss_value = 0
for i in tqdm.tqdm(range(step_num)):
x _batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
points _batch, mask_batch = get_points()
# print(x_batch.shape)
# print(mask_batch.shape)
dg _loss_n = exe.run(dg_program,
feed={'x': x_batch,
'mask':mask_batch,},
fetch _list=[dg_loss])[0]
g _loss_value += dg _loss_n
print('Pass _id:{}, Completion loss: {}'.format(pass_id, g _loss_value))
np.random.shuffle(x_test)
x _batch = x_test[:BATCH_SIZE]
completion _n = exe.run(dg_program,
feed={'x': x_batch,
'mask': mask_batch,},
fetch_list=[ completion])[ 0][ 0]
# 修复图片
sample = np.array((completion_n + 1) * 127.5, dtype=np.uint8)
# 原图
x _im = np.array((x_batch[0] + 1) * 127.5, dtype=np.uint8)
# 挖空洞输入图
input _im_data = x _im * (1 - mask_batch[0])
input _im = np.array(input_im _data + np.ones_like(x _im) * mask_batch[0] * 255, dtype=np.uint8)
output _im = np.concatenate((x_im,input_im,sample),axis=1)
#print(output_im.shape)
cv2.imwrite('./output/pass _id:{}.jpg'.format(pass_id), cv2.cvtColor(output _im, cv2.COLOR_RGB2BGR))
# 保存模型
save _pretrain_model_path = 'models/'
# 创建保持模型文件目录
#os.makedirs(save _pretrain_model_path)
fluid.io.save _params(executor=exe, dirname=save_pretrain _model_path, main _program=dg_program)
# 生成器判断器一起训练
else:
g _loss_value = 0
d _loss_value = 0
for i in tqdm.tqdm(range(step_num)):
x _batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
points _batch, mask_batch = get_points()
dg _loss_n = exe.run(dg_program,
feed={'x': x_batch,
'mask':mask_batch,},
fetch _list=[dg_loss])[0]
g _loss_value += dg _loss_n
completion _n = exe.run(dg_program,
feed={'x': x_batch,
'mask': mask_batch,},
fetch_list=[completion])[0]
local _x_batch = []
local _completion_batch = []
for i in range(BATCH_SIZE):
x1, y1, x2, y2 = points_batch[i]
local _x_batch.append(x_batch[ i][ y1:y2, x1:x2, :])
local _completion_batch.append(completion_n[ i][ y1:y2, x1:x2, :])
local _x_batch = np.array(local _x_batch)
local _completion_batch = np.array(local _completion_batch)
d _loss_n = exe.run(d_program,
feed={'x': x _batch, 'mask': mask_batch, 'local _x': local_x _batch, 'global_completion': completion _n, 'local_completion': local _completion_batch},
fetch _list=[d_loss])[0]
d _loss_value += d _loss_n
print('Pass _id:{}, Completion loss: {}'.format(pass_id, g _loss_value))
print('Pass _id:{}, Discriminator loss: {}'.format(pass_id, d _loss_value))
np.random.shuffle(x_test)
x _batch = x_test[:BATCH_SIZE]
completion _n = exe.run(dg_program,
feed={'x': x_batch,
'mask': mask_batch,},
fetch_list=[ completion])[ 0][ 0]
# 修复图片
sample = np.array((completion_n + 1) * 127.5, dtype=np.uint8)
# 原图
x _im = np.array((x_batch[0] + 1) * 127.5, dtype=np.uint8)
# 挖空洞输入图
input _im_data = x _im * (1 - mask_batch[0])
input _im = np.array(input_im _data + np.ones_like(x _im) * mask_batch[0] * 255, dtype=np.uint8)
output _im = np.concatenate((x_im,input_im,sample),axis=1)
#print(output_im.shape)
cv2.imwrite('./output/pass _id:{}.jpg'.format(pass_id), cv2.cvtColor(output _im, cv2.COLOR_RGB2BGR))
# 保存模型
save _pretrain_model_path = 'models/'
# 创建保持模型文件目录
#os.makedirs(save _pretrain_model_path)
fluid.io.save _params(executor=exe, dirname=save_pretrain _model_path, main _program = dg_program)