"""Helper functions for parsing caffe prototxt into a workable DAG """ def process_network_proto(caffe_root, deploy_proto): """ Runs the caffe upgrade tool on the prototxt to create a prototxt in the latest format. This enable us to work just with latest structures, instead of supporting all the variants :param caffe_root: link to caffe root folder, where the upgrade tool is located :param deploy_proto: name of the original prototxt file :return: name of new processed prototxt file """ processed_deploy_proto = deploy_proto + ".processed" from shutil import copyfile copyfile(deploy_proto, processed_deploy_proto) # run upgrade tool on new file name (same output file) import os upgrade_tool_command_line = caffe_root + '/build/tools/upgrade_net_proto_text.bin ' \ + processed_deploy_proto + ' ' + processed_deploy_proto os.system(upgrade_tool_command_line) return processed_deploy_proto class LayerRecord(object): """ A record which describe basic layer parameters """ def __init__(self, layer_def): self.layer_def = layer_def self.name = layer_def.name self.type = layer_def.type # keep filter, stride and pad if layer_def.type == 'Convolution': if LayerRecord._is_iterable(layer_def.convolution_param.kernel_size): self.filter = list(layer_def.convolution_param.kernel_size) else: self.filter = list([layer_def.convolution_param.kernel_size]) if len(self.filter) == 1: self.filter *= 2 if LayerRecord._is_iterable(layer_def.convolution_param.pad): self.pad = list(layer_def.convolution_param.pad) else: self.pad = list([layer_def.convolution_param.pad]) if len(self.pad) == 0: self.pad = [0, 0] elif len(self.pad) == 1: self.pad *= 2 if LayerRecord._is_iterable(layer_def.convolution_param.stride): self.stride = list(layer_def.convolution_param.stride) else: self.stride = list([layer_def.convolution_param.stride]) if len(self.stride) == 0: self.stride = [1, 1] elif len(self.stride) == 1: self.stride *= 2 elif layer_def.type == 'Pooling': self.filter = [layer_def.pooling_param.kernel_size] if len(self.filter) == 1: self.filter *= 2 self.pad = [layer_def.pooling_param.pad] if len(self.pad) == 0: self.pad = [0, 0] elif len(self.pad) == 1: self.pad *= 2 self.stride = [layer_def.pooling_param.stride] if len(self.stride) == 0: self.stride = [1, 1] elif len(self.stride) == 1: self.stride *= 2 else: self.filter = [0, 0] self.pad = [0, 0] self.stride = [1, 1] # keep tops self.tops = list(layer_def.top) # keep bottoms self.bottoms = list(layer_def.bottom) # list of parent layers self.parents = [] # list of child layers self.children = [] @staticmethod def _is_iterable(obj): return hasattr(obj, '__iter__') def read_network_dag(processed_deploy_prototxt): """ Reads from the caffe prototxt the network structure :param processed_deploy_prototxt: name of prototxt to load, preferably the prototxt should be processed before using a call to process_network_proto() :return: network_def, layer_name_to_record, top_to_layers network_def: caffe network structure, gives access to *all* the network information layer_name_to_record: *ordered* dictionary which maps between layer name and a structure which describes in a simple form the layer parameters top_to_layers: dictionary which maps a blob name to an ordered list of layers which output it when a top is used several times, like in inplace layhers, the list will contain all the layers by order of appearance """ from caffe.proto import caffe_pb2 from google.protobuf import text_format from collections import OrderedDict # load prototxt file network_def = caffe_pb2.NetParameter() with open(processed_deploy_prototxt, 'r') as proto_file: text_format.Merge(str(proto_file.read()), network_def) # map layer name to layer record layer_name_to_record = OrderedDict() for layer_def in network_def.layer: if (len(layer_def.include) == 0) or \ (caffe_pb2.TEST in [item.phase for item in layer_def.include]): layer_name_to_record[layer_def.name] = LayerRecord(layer_def) top_to_layers = dict() for layer in network_def.layer: # no specific phase, or TEST phase is specifically asked for if (len(layer.include) == 0) or (caffe_pb2.TEST in [item.phase for item in layer.include]): for top in layer.top: if top not in top_to_layers: top_to_layers[top] = list() top_to_layers[top].append(layer.name) # find parents and children of all layers for child_layer_name in layer_name_to_record.keys(): # pylint: disable=too-many-nested-blocks child_layer_def = layer_name_to_record[child_layer_name] for bottom in child_layer_def.bottoms: if bottom in top_to_layers: for parent_layer_name in top_to_layers[bottom]: if parent_layer_name in layer_name_to_record: parent_layer_def = layer_name_to_record[parent_layer_name] if parent_layer_def not in child_layer_def.parents: child_layer_def.parents.append(parent_layer_def) if child_layer_def not in parent_layer_def.children: parent_layer_def.children.append(child_layer_def) # update filter, strid, pad for maxout "structures" for layer_name in layer_name_to_record.keys(): layer_def = layer_name_to_record[layer_name] if layer_def.type == 'Eltwise' and \ len(layer_def.parents) == 1 and \ layer_def.parents[0].type == 'Slice' and \ len(layer_def.parents[0].parents) == 1 and \ layer_def.parents[0].parents[0].type in ['Convolution', 'InnerProduct']: layer_def.filter = layer_def.parents[0].parents[0].filter layer_def.stride = layer_def.parents[0].parents[0].stride layer_def.pad = layer_def.parents[0].parents[0].pad return network_def, layer_name_to_record, top_to_layers def read_caffe_mean(caffe_mean_file): """ Reads caffe formatted mean file :param caffe_mean_file: path to caffe mean file, presumably with 'binaryproto' suffix :return: mean image, converted from BGR to RGB format """ import caffe_parser import numpy as np mean_blob = caffe_parser.caffe_pb2.BlobProto() with open(caffe_mean_file, 'rb') as f: mean_blob.ParseFromString(f.read()) img_mean_np = np.array(mean_blob.data) img_mean_np = img_mean_np.reshape(mean_blob.channels, mean_blob.height, mean_blob.width) # swap channels from Caffe BGR to RGB img_mean_np[[0, 2], :, :] = img_mean_np[[2, 0], :, :] return img_mean_np