ssd.py 文件源码

def __init__(self):
        # load MS COCO labels
        labelmap_file = os.path.join(CAFFE_ROOT, LABEL_MAP)
        file = open(labelmap_file, 'r')
        self._labelmap = caffe_pb2.LabelMap()
        text_format.Merge(str(file.read()), self._labelmap)

        model_def = os.path.join(CAFFE_ROOT, PROTO_TXT)
        model_weights = os.path.join(CAFFE_ROOT, CAFFE_MODEL)

        self._net = caffe.Net(model_def, model_weights, caffe.TEST)
        self._transformer = caffe.io.Transformer(
            {'data': self._net.blobs['data'].data.shape})
        self._transformer.set_transpose('data', (2, 0, 1))
        self._transformer.set_mean('data', np.array([104, 117, 123]))
        self._transformer.set_raw_scale('data', 255)
        self._transformer.set_channel_swap('data', (2, 1, 0))

        # set net to batch size of 1
        image_resize = IMAGE_SIZE
        self._net.blobs['data'].reshape(1, 3, image_resize, image_resize)