手把手教物体检测——RFBNet

1. 下载代码

https://github.com/ruinmessi/RFBNet

1. 解压后在RFBNet的data文件夹下新建VOCdevkit文件夹，将VOC数据集放进去。
2. 修改类别。

在voc0712.py中的VOC_CLASSES中的类别修改为自己数据集的类别。修改后：

VOC_CLASSES = ( '__background__', # always index 0

   'aircraft', 'oiltank')

注意：第一个类别是背景，不用修改。

修改config.py的文件路径。

RBFNet默认的路径是linux的路径，我使用的是Win10，需要修改路径，否则找不到数据集。

将：

# gets home dir cross platform
home = os.path.expanduser("~")
ddir = os.path.join(home,"data/VOCdevkit/")
# note: if you used our download scripts, this should be right
VOCroot = ddir # path to VOCdevkit root dir
COCOroot = os.path.join(home,"data/COCO/")

改为：

# gets home dir cross platform
ddir = "data/VOCdevkit/"
# note: if you used our download scripts, this should be right
VOCroot = ddir # path to VOCdevkit root dir
COCOroot = "data/COCO/"

1. 修改utils->nms_wrapper.py

这个文件的作用的调用nms中文件，nms指的是非极大值抑制。

nms文件夹是集中nms编写的方式，采用py的即可，性能上不会有太大的影响。

将：

from .nms.cpu_nms import cpu_nms, cpu_soft_nms
from .nms.gpu_nms import gpu_nms

修改为：

from .nms.py_cpu_nms import py_cpu_nms

将：

def nms(dets, thresh, force_cpu=False):
    """Dispatch to either CPU or GPU NMS implementations."""
    if dets.shape[0] == 0:
        return []
    if force_cpu:
        #return cpu_soft_nms(dets, thresh, method = 0)
        return cpu_nms(dets, thresh)
    return gpu_nms(dets, thresh)

修改为：

def nms(dets, thresh, force_cpu=False):
    """Dispatch to either CPU or GPU NMS implementations."""
    if dets.shape[0] == 0:
        return []
    if force_cpu:
        #return cpu_soft_nms(dets, thresh, method = 0)
        return py_cpu_nms(dets, thresh)
    return py_cpu_nms(dets, thresh)

新建weights文件，下载vgg16模型放到里面。

下载地址：https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

修改data->coco.py

将：

from utils.pycocotools.coco import COCO

from utils.pycocotools.cocoeval import COCOeval

from utils.pycocotools import mask as COCOmask

修改为：

from pycocotools.coco import COCO

from pycocotools.cocoeval import COCOeval

from pycocotools import mask as COCOmask

删除utils->pycocotools文件夹 。

修改train_RFB.py

修改全局参数：

parser = argparse.ArgumentParser(

   description='Receptive Field Block Net Training')

parser.add_argument('-v', '--version', default='RFB_vgg',

                   help='RFB_vgg ,RFB_E_vgg or RFB_mobile version.')

parser.add_argument('-s', '--size', default='512',

                   help='300 or 512 input size.')

parser.add_argument('-d', '--dataset', default='VOC',

                   help='VOC or COCO dataset')

parser.add_argument(

   '--basenet', default='./weights/vgg16_reducedfc.pth', help='pretrained base model')

parser.add_argument('--jaccard_threshold', default=0.5,

                   type=float, help='Min Jaccard index for matching')

parser.add_argument('-b', '--batch_size', default=2,

                   type=int, help='Batch size for training')

parser.add_argument('--num_workers', default=2,

                   type=int, help='Number of workers used in dataloading')

parser.add_argument('--cuda', default=True,

                   type=bool, help='Use cuda to train model')

parser.add_argument('--ngpu', default=1, type=int, help='gpus')

parser.add_argument('--lr', '--learning-rate',

                   default=4e-3, type=float, help='initial learning rate')

parser.add_argument('--momentum', default=0.9, type=float, help='momentum')

parser.add_argument(

   '--resume_net', default=None, help='resume net for retraining')

parser.add_argument('--resume_epoch', default=0,

                   type=int, help='resume iter for retraining')

parser.add_argument('-max','--max_epoch', default=300,

                   type=int, help='max epoch for retraining')

parser.add_argument('--weight_decay', default=5e-4,

                   type=float, help='Weight decay for SGD')

parser.add_argument('--gamma', default=0.1,

                   type=float, help='Gamma update for SGD')

parser.add_argument('--log_iters', default=True,

                   type=bool, help='Print the loss at each iteration')

parser.add_argument('--save_folder', default='./weights/',

                   help='Location to save checkpoint models')

将：

if args.dataset == 'VOC':
    train_sets = [('2007', 'trainval'), ('2012', 'trainval')]
    cfg = (VOC_300, VOC_512)[args.size == '512']
else:
    train_sets = [('2014', 'train'),('2014', 'valminusminival')]
    cfg = (COCO_300, COCO_512)[args.size == '512']

修改为：

if args.dataset == 'VOC':
    train_sets = [('2007', 'trainval'), 
    cfg = (VOC_300, VOC_512)[args.size == '512']
else:
    train_sets = [('2014', 'train'),('2014', 'valminusminival')]
    cfg = (COCO_300, COCO_512)[args.size == '512']

将82行：

num_classes = (21, 81)[args.dataset == 'COCO']

修改为：

num_classes = (3, 81)[args.dataset == 'COCO']#如果是COCO就选择81,3是本次的类别+1（背景）

结果：

前5个Epoch将学习率从小升到初始值，是用来对模型进行热身。

测试，并验证测试结果。

修改test_RFB.py

修改全局参数

parser.add_argument('-v', '--version', default='RFB_vgg',

                   help='RFB_vgg ,RFB_E_vgg or RFB_mobile version.')#和训练的模型保持一致。

parser.add_argument('-s', '--size', default='512',

                   help='300 or 512 input size.')#和训练是选用的大小保持一致。

parser.add_argument('-d', '--dataset', default='VOC',

                   help='VOC or COCO version')

parser.add_argument('-m', '--trained_model', default='weights/Final_RFB_vgg_VOC.pth',

                   type=str, help='Trained state_dict file path to open')#选择训练好的模型

parser.add_argument('--cuda', default=False, type=bool,

                   help='Use cuda to train model')

parser.add_argument('--cpu', default=True, type=bool,

                   help='Use cpu nms')

parser.add_argument('--retest', default=False, type=bool,

                   help='test cache results')

args = parser.parse_args()

将148行：

                  num_classes = (21, 81)[args.dataset == 'COCO']

修改为：

                  num_classes = (3, 81)[args.dataset == 'COCO']

将71行：

                  num_classes = (21, 81)[args.dataset == 'COCO']

修改为：

                  num_classes = (3, 81)[args.dataset == 'COCO']

修改voc0712.py的281行

将：

                  annopath = os.path.join(

                       rootpath,

                       'Annotations',

                       '{:s}.xml')

修改为：annopath = rootpath+'/Annotations/{:s}.xml'#解决验证时找不到测试集xml的问题。

运行test_RFB.py结果如下：

测试单张图片，并展示结果。

from __future__ import print_function

import torch

import torch.backends.cudnn as cudnn

import os

import argparse

import numpy as np

from matplotlib import pyplot as plt

from data import AnnotationTransform, COCODetection, VOCDetection, BaseTransform, VOC_300, VOC_512, COCO_300, COCO_512, \

   COCO_mobile_300

from layers.functions import Detect, PriorBox

from utils.nms_wrapper import nms

import cv2

from data import VOC_CLASSES as labels

from collections import OrderedDict

import time

#功能：测试单一的一张图片

parser = argparse.ArgumentParser(description='Receptive Field Block Net')

parser.add_argument('-v', '--version', default='RFB_vgg',

                   help='RFB_vgg ,RFB_E_vgg or RFB_mobile version.')

parser.add_argument('-s', '--size', default='512',

                   help='300 or 512 input size.')

parser.add_argument('-n', '--num_classes', default='3',

                   help='300 or 512 input size.')

parser.add_argument('-d', '--dataset', default='VOC',

                   help='VOC or COCO version')

parser.add_argument('-m', '--trained_model', default='weights/RFB_vgg_VOC_epoches_160.pth',

                   type=str, help='Trained state_dict file path to open')

parser.add_argument('--save_folder', default='eval/', type=str,

                   help='Dir to save results')

parser.add_argument('--cuda', default=True, type=bool,

                   help='Use cuda to train model')

parser.add_argument('--cpu', default=False, type=bool,

                   help='Use cpu nms')

parser.add_argument('--retest', default=False, type=bool,

                   help='test cache results')

args = parser.parse_args()

if not os.path.exists(args.save_folder):

   os.mkdir(args.save_folder)

if args.dataset == 'VOC':

   cfg = (VOC_300, VOC_512)[args.size == '512']

else:

   cfg = (COCO_300, COCO_512)[args.size == '512']

if args.version == 'RFB_vgg':

   from models.RFB_Net_vgg import build_net

elif args.version == 'RFB_E_vgg':

   from models.RFB_Net_E_vgg import build_net

elif args.version == 'RFB_mobile':

   from models.RFB_Net_mobile import build_net

   cfg = COCO_mobile_300

else:

   print('Unkown version!')

priorbox = PriorBox(cfg)

with torch.no_grad():

   priors = priorbox.forward()

   if args.cuda:

       priors = priors.cuda()

t1=time.time()

imagePath = "data/VOCdevkit/aircraft_27.jpg"

# load net

img_dim = int(args.size)

num_classes = int(args.num_classes)

net = build_net('test', img_dim, num_classes)  # initialize detector

state_dict = torch.load(args.trained_model)

# create new OrderedDict that does not contain `module.`

new_state_dict = OrderedDict()

for k, v in state_dict.items():

   head = k[:7]

   if head == 'module.':

       name = k[7:]  # remove `module.`

   else:

       name = k

   new_state_dict[name] = v

net.load_state_dict(new_state_dict)

net.eval()

print('Finished loading model!')

if args.cuda:

   net = net.cuda()

   cudnn.benchmark = True

else:

   net = net.cpu()

top_k = 200

detector = Detect(num_classes, 0, cfg)

save_folder = os.path.join(args.save_folder, args.dataset)

if not os.path.exists(save_folder):

   os.mkdir(save_folder)

# dump predictions and assoc. ground truth to text file for now

det_file = os.path.join(save_folder, 'detections.pkl')

image = cv2.imread(imagePath, cv2.IMREAD_COLOR)

rgb_means = ((104, 117, 123), (103.94, 116.78, 123.68))[args.version == 'RFB_mobile']

scale = torch.Tensor([image.shape[1], image.shape[0],

                     image.shape[1], image.shape[0]])

transform = BaseTransform(net.size, rgb_means, (2, 0, 1))

with torch.no_grad():

   x = transform(image).unsqueeze(0)

   if args.cuda:

       x = x.cuda()

       scale = scale.cuda()

out = net(x)  # forward pass

boxes, scores = detector.forward(out, priors)

boxes = boxes[0]

scores = scores[0]

boxes *= scale

boxes = boxes.cpu().numpy()

scores = scores.cpu().numpy()

result = []

for j in range(1, num_classes):

   inds = np.where(scores[:, j] > 0.99)[0]

   if len(inds) == 0:

       continue

   label_name = labels[j]

   c_bboxes = boxes[inds]

   c_scores = scores[inds, j]

   c_dets = np.hstack((c_bboxes, c_scores[:, np.newaxis])).astype(

       np.float32, copy=False)

   keep = nms(c_dets, 0.45, force_cpu=args.cpu)

   c_dets = c_dets[keep, :]

   for listbox in c_dets:

       temp = []

       temp.append(label_name)

       temp.append(listbox[4])

       temp.append(int(listbox[0]))

       temp.append(int(listbox[1]))

       temp.append(int(listbox[2]))

       temp.append(int(listbox[3]))

       result.append(temp)

print(result)

t2=time.time()

print(t2-t1)

isShowResult = True

if isShowResult:

   plt.figure(figsize=(10, 10))

   colors = plt.cm.hsv(np.linspace(0, 1, num_classes)).tolist()

   rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

   plt.imshow(rgb_image)  # plot the image for matplotlib

   currentAxis = plt.gca()

   for listbox in result:

       label_name = listbox[0]

       i = labels.index(label_name)

       score = listbox[1]

       coords = (listbox[2], listbox[3]), listbox[4] - listbox[2] + 1, listbox[5] - listbox[3] + 1

       display_txt = '%s: %.2f' % (label_name, score)

       color = colors[i]

       currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))

       currentAxis.text(listbox[2], listbox[3], display_txt, bbox={'facecolor': color, 'alpha': 0.5})

   plt.show(