pvanet 修改demo.py来载入自己训练过的模型并显示检测结果

在CSDN上找相关的博文，都没有找到我想要的。

我的目的是，已经用pvanet训练好了符合自己需求的模型。怎么修改demo.py载入自己的模型来显示检测结果。

先说下我的情况：我根据自己的需求，用pvanet这个网络结构，预训练模型采用test.model或者original.model来训练的，目的是识别数字。

想改动pvanet自带的tool文件夹下的demo.py（与py-fasterrcnn的demo.py一样的）来测试我的模型检测效果。

demo.py需要改动的地方：

1、开始的地方修改类别：

CLASSES = ('__background__',  
           '0', '1', '2', '3',  
           '4', '5', '6', '7', '8',  
           '9')  

2、修改你prototxt的路径:

把这一句

prototxt = os.path.join(cfg.MODELS_DIR, NETS[args.demo_net][0],  
                            'faster_rcnn_alt_opt', 'faster_rcnn_test.pt')  

改成

prototxt = os.path.join('$pva-faster-rcnn/models/pvanet/example_train_384/test.prototxt')  //test.prototxt是你训练模型时对应网络结构的test网络结构  

3、修改你caffemodel的路径：

把这一句

caffemodel = os.path.join(cfg.DATA_DIR, 'faster_rcnn_models',  
                              NETS[args.demo_net][1])  

改成

caffemodel = os.path.join('$pva-faster-rcnn/output/faster_rcnn_pvanet/voc_2007_trainval/pvanet_frcnn_384_iter_30000.caffemodel')  

4、修改需要测试图片的路径

im_names = ['000456.jpg', '000542.jpg', '001150.jpg',  
                '001763.jpg', '004545.jpg','timg1.jpg']  
    for im_name in im_names:  
        print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'  
        print 'Demo for data/demo/{}'.format(im_name)  //把图片换成你想要测试的图片名字，在data/demo/下  

这些图片的名字要与文件夹内的图片对应的上，如果这边传入文件名，文件夹下没有相关图片，会报错。

关键的地方来了，按照以上四步都修改完，会报错：

F0808 15:20:33.498361 22337 concat_layer.cpp:42] Check failed: top_shape[j] == bottom[i]->shape(j) (38 vs. 37) All inputs must have the same shape, except at concat_axis.
*** Check failure stack trace: ***
Aborted (core dumped)

Git上查了一下原因，需要传cfg的参数。所以又有了下面的修改

5、关于cfg参数的修改

首先在开头部分：

from fast_rcnn.config import cfg, cfg_from_file ///后面要加上, cfg_from_file 

然后在选择参数部分：

def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description='Faster R-CNN demo')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]',
                        default=0, type=int)
    parser.add_argument('--cpu', dest='cpu_mode',
                        help='Use CPU mode (overrides --gpu)',
                        action='store_true')
    parser.add_argument('--net', dest='demo_net', help='Network to use [vgg16]',
                        choices=NETS.keys(), default='vgg16')
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file', default=None, type=str)
    args = parser.parse_args()

    return args

相比较原demo.py这里要加上传入cfg参数的代码

parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file', default=None, type=str)

然后在这里修改：

if not os.path.isfile(caffemodel):
        raise IOError(('{:s} not found.\nDid you run ./data/script/'
                       'fetch_faster_rcnn_models.sh?').format(caffemodel))
    if args.cfg_file is not None:
        cfg_from_file(args.cfg_file)

    if args.cpu_mode:
        caffe.set_mode_cpu()
    else:
        caffe.set_mode_gpu()
        caffe.set_device(args.gpu_id)
        cfg.GPU_ID = args.gpu_id
    net = caffe.Net(prototxt, caffemodel, caffe.TEST)

    print '\n\nLoaded network {:s}'.format(caffemodel)

相比较原demo.py这里加一个判断，加上了这一段代码，判断是否cfg文件是否传入

if args.cfg_file is not None:
        cfg_from_file(args.cfg_file)

全部修改完以后，保存脚本，运行脚本，并传入参数

./tools/demo.py --cfg  models/pvanet/cfgs/submit_0716.yml