| """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 |
