mtcnn_aligner.py 文件源码

def get_aligner(caffe_model_path, use_more_stage=False):
    caffe.set_mode_gpu()
#    PNet = caffe.Net(caffe_model_path + "/det1.prototxt",
#                     caffe_model_path + "/det1.caffemodel", caffe.TEST)
    if use_more_stage:
        RNet = caffe.Net(caffe_model_path + "/det2.prototxt",
                         caffe_model_path + "/det2.caffemodel", caffe.TEST)
    else:
        RNet = None

    ONet = caffe.Net(caffe_model_path + "/det3.prototxt",
                     caffe_model_path + "/det3.caffemodel", caffe.TEST)

    LNet = caffe.Net(caffe_model_path + "/det4.prototxt",
                     caffe_model_path + "/det4.caffemodel", caffe.TEST)

#    return (PNet, RNet, ONet)
    return (RNet, ONet, LNet)
#    return (RNet, ONet, None)