#!/usr/bin/env?python
#?-*-?coding:?UTF-8?-*-

##?This?tool?is?used?to?merge?‘Conv-BN-Scale‘?into?a?single?‘Conv‘?layer.?
##?It?is?copied?from?https://github.com/sanghoon/pva-faster-rcnn/blob/master/tools/gen_merged_model.py

import?numpy?as?np
import?sys
import?os
import?os.path?as?osp
import?google.protobuf?as?pb
import?google.protobuf.text_format
from?argparse?import?ArgumentParser
import?caffe
#sys.path.append（“./python“）
#caffe_root?=?‘E:/caffe-ssd/caffe-microsoft/Build/x64/Release/‘
#sys.path.insert（0?caffe_root?+?‘pycaffe‘）


caffe.set_mode_cpu（）


def?load_and_fill_biases（src_model?src_weights?dst_model?dst_weights）:
????with?open（src_model）?as?f:
????????model?=?caffe.proto.caffe_pb2.NetParameter（）
????????pb.text_format.Merge（f.read（）?model）

????for?i?layer?in?enumerate（model.layer）:
????????if?layer.type?==?‘Convolution‘:?#?or?layer.type?==?‘Scale‘:
????????????#?Add?bias?layer?if?needed
????????????if?layer.convolution_param.bias_term?==?False:
????????????????layer.convolution_param.bias_term?=?True
????????????????layer.convolution_param.bias_filler.type?=?‘constant‘
????????????????layer.convolution_param.bias_filler.value?=?0.0

????with?open（dst_model?‘w‘）?as?f:
????????f.write（pb.text_format.MessageToString（model））

????caffe.set_mode_cpu（）
????net_src?=?caffe.Net（src_model?src_weights?caffe.TEST）
????net_dst?=?caffe.Net（dst_model?caffe.TEST）
????for?key?in?net_src.params.keys（）:
????????for?i?in?range（len（net_src.params[key]））:
????????????net_dst.params[key][i].data[:]?=?net_src.params[key][i].data[:]

????if?dst_weights?is?not?None:
????????#?Store?params
????????pass

????return?net_dst


def?merge_conv_and_bn（net?i_conv?i_bn?i_scale）:
????#?This?is?based?on?Kyeheyon‘s?work
????assert（i_conv?!=?None）
????assert（i_bn?!=?None）

????def?copy_double（data）:
????????return?np.array（data?copy=True?dtype=np.double）

????key_conv?=?net._layer_names[i_conv]
????key_bn?=?net._layer_names[i_bn]
????key_scale?=?net._layer_names[i_scale]?if?i_scale?else?None

????#?Copy
????bn_mean?=?copy_double（net.params[key_bn][0].data）
????bn_variance?=?copy_double（net.params[key_bn][1].data）
????num_bn_samples?=?copy_double（net.params[key_bn][2].data）

????#?and?Invalidate?the?BN?layer
????net.params[key_bn][0].data[:]?=?0
????net.params[key_bn][1].data[:]?=?1
????net.params[key_bn][2].data[:]?=?1

????if?num_bn_samples[0]?==?0:
????????num_bn_samples[0]?=?1

????if?net.params.has_key（key_scale）:
????????print?‘Combine?{:s}?+?{:s}?+?{:s}‘.format（key_conv?key_bn?key_scale）
????????scale_weight?=?copy_double（net.params[key_scale][0].data）
????????scale_bias?=?copy_double（net.params[key_scale][1].data）
????????net.params[key_scale][0].data[:]?=?1
????????net.params[key_scale][1].data[:]?=?0

????else:
????????print?‘Combine?{:s}?+?{:s}‘.format（key_conv?key_bn）
????????scale_weight?=?1
????????scale_bias?=?0

????weight?=?copy_double（net.params[key_conv][0].data）
????bias?=?copy_double（net.params[key_conv][1].data）

????alpha?=?scale_w