demo.py 文件源码

def main(argv):
    model_filename = ''
    weight_filename = ''
    img_filename = ''
    try:
        opts, args = getopt.getopt(argv, "hm:w:i:")
        print opts
    except getopt.GetoptError:
        print 'yolo_main.py -m <model_file> -w <output_file> -i <img_file>'
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print 'yolo_main.py -m <model_file> -w <weight_file> -i <img_file>'
            sys.exit()
        elif opt == "-m":
            model_filename = arg
        elif opt == "-w":
            weight_filename = arg
        elif opt == "-i":
            img_filename = arg
    print 'model file is "', model_filename
    print 'weight file is "', weight_filename
    print 'image file is "', img_filename

    caffe.set_device(0)
    caffe.set_mode_gpu()
    net = caffe.Net(model_filename, weight_filename, caffe.TEST)
    img = caffe.io.load_image(img_filename)  # load the image using caffe io
    img_ = scipy.misc.imresize(img, (448, 448))
    transformer = SimpleTransformer([104.00699, 116.66877, 122.67892])
    input = transformer.preprocess(img_)
    out = net.forward_all(data=input)
    print out.iteritems()
    img_cv = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    results = interpret_output(out['result'][0], img.shape[1], img.shape[0])  # fc27 instead of fc12 for yolo_small
    show_results(img_cv, results, img.shape[1], img.shape[0])
    cv2.waitKey(0)