Faster R-CNN 修改demo.py为批量测试

最新推荐文章于 2021-03-10 13:17:37 发布

南石北岸生

最新推荐文章于 2021-03-10 13:17:37 发布

阅读量3.3k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： Faster R-CNN 文章标签： Python

本文链接：https://blog.youkuaiyun.com/gusui7202/article/details/83239142

本文档介绍了如何修改Faster R-CNN的demo.py脚本以实现批量图片测试。主要修改包括从txt文件读取图片列表、手动设置模型路径以及添加保存检测结果图片的功能。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

修改了tools/demo.py以下几点：

1.加入：读取txt文档内的图片名称合集为列表。

2.模型的路径改为手动选择

3.添加保存图片函数和路径。

使用：

修改类别、修改名称、修改模型地址、修改图片保存地址、修改要读入的txt地址。

#!/usr/bin/env python

"""
https://blog.youkuaiyun.com/gusui7202/article/details/83239142
qhy。
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths
from model.config import cfg
from model.test import im_detect
from model.nms_wrapper import nms

from utils.timer import Timer
import tensorflow as tf
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import os, cv2
import argparse


from nets.vgg16 import vgg16
from nets.resnet_v1 import resnetv1

CLASSES = ('__background__',  # always index 0
                     'normal bolt','normal bolt-2','shim losing','nut losing','nut directly loosening','pin loosening','pin closing','visible pin losing','invisible pin losing')#修改类别

NETS = {'vgg16': ('vgg16_faster_rcnn_iter_70000.ckpt',),'res101': ('res101_faster_rcnn_iter_70000.ckpt',)}
DATASETS= {'pascal_voc': ('voc_2007_trainval',),'pascal_voc_0712': ('voc_2007_trainval+voc_2012_trainval',)}#这里要修改一下名字，与要测试的模型名字一致

def vis_detections(image_name, im, class_name, dets, thresh=0.5):
    """Draw detected bounding boxes."""
    inds = np.where(dets[:, -1] >= thresh)[0]
    if len(inds) == 0:
        return

    im = im[:, :, (2,