def __init__(self, model_file, pretrained_file, mean_value=None,
layer=['pool5'], input_size = None ):
caffe.set_mode_gpu()
caffe.Net.__init__(self, model_file, pretrained_file, caffe.TEST)
# get name input layer
self.list_layers = layer
self.mean_value = mean_value
# set transformer object
self.transformer = caffe.io.Transformer({'data': self.blobs['data'].data.shape})
self.transformer.set_transpose( 'data', (2,0,1) )
if mean_value is not None:
self.transformer.set_mean('data', mean_value)
self.transformer.set_raw_scale('data', 255)
self.transformer.set_channel_swap('data', (2,1,0))
if input_size is not None:
#reshape the input
print "New input! {}".format(input_size)
self.reshape_input( input_size[0], input_size[1], input_size[2], input_size[3] )
评论列表
文章目录