- скачать код
- После распаковки создайте новую папку VOCdevkit в папке данных RFBNet и поместите в нее набор данных VOC.
- Изменить категорию.
Категория в VOC_CLASSES в voc0712.py изменена на категорию собственного набора данных. После модификации:
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/"
- Изменить 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)
- Создайте новый файл весов, скачайте модель vgg16 и поместите ее в него.
ссылка для скачивания: Да 3. Amazon AWS.com/AMD, например, root — нет...
- Изменить данные->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 эпох увеличивают скорость обучения от небольшой до начального значения, которое используется для прогрева модели.
- проверить и проверить результаты теста.
Изменить 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' ]
Измените строку 281 файла voc0712.py.
будет:
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()