手把手教物体检测——RFBNet

举报
AI浩 发表于 2021/12/23 01:31:37 2021/12/23
【摘要】 下载代码 https://github.com/ruinmessi/RFBNet 解压后在RFBNet的data文件夹下新建VOCdevkit文件夹,将VOC数据集放进去。修改类别。 在voc0712.py中的VOC_CLASSES中的类别修改为自己数据集的类别。修改后: VOC_CLASSES = ( '__backgro...
  1. 下载代码

https://github.com/ruinmessi/RFBNet

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

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

VOC_CLASSES = ( '__background__', # always index 0
   
'aircraft', 'oiltank')

注意:第一个类别是背景,不用修改。

  1. 修改config.py的文件路径。

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

将:


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

改为:


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

  1. 修改utils->nms_wrapper.py

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

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

将:


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

修改为:

from .nms.py_cpu_nms import py_cpu_nms
 

将:


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

修改为:


  
  1. def nms(dets, thresh, force_cpu=False):
  2.     """Dispatch to either CPU or GPU NMS implementations."""
  3.     if dets.shape[0] == 0:
  4.         return []
  5.     if force_cpu:
  6.         #return cpu_soft_nms(dets, thresh, method = 0)
  7.         return py_cpu_nms(dets, thresh)
  8.     return py_cpu_nms(dets, thresh)
  1. 新建weights文件,下载vgg16模型放到里面。

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

  1. 修改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文件夹 。

  1. 修改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')

将:


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

修改为:


  
  1. if args.dataset == 'VOC':
  2.     train_sets = [('2007', 'trainval'),
  3.     cfg = (VOC_300, VOC_512)[args.size == '512']
  4. else:
  5.     train_sets = [('2014', 'train'),('2014', 'valminusminival')]
  6.     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将学习率从小升到初始值,是用来对模型进行热身。

  1. 测试,并验证测试结果。

修改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结果如下:

  1. 测试单张图片,并展示结果。

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()

文章来源: wanghao.blog.csdn.net,作者:AI浩,版权归原作者所有,如需转载,请联系作者。

原文链接:wanghao.blog.csdn.net/article/details/105876651

【版权声明】本文为华为云社区用户转载文章,如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。