Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
(0)环境
Windows10
PyTorch1.4
pycharm
(一)准备数据集
首先新建了VOC2007文件夹,然后再在里面新建三个文件夹,如下图。
其中,JPEGImages用来存放照片,Annotations存放xml文件。
第一步:对图片进行标注。我这里选了五种类别(书包,钟表,手机,水杯,鼠标)各十张图片,总共五十张照片(由于显卡不行,只选了这么多),标注工具为labelImg,教程网上很多,这里不再赘述。标注得到的xml文件存放在Annotations中。
第二步:用make_txt.py文件将数据集分为4个txt文件,用于后续的训练与测试。下图为make_txt.py文件以及分好的四个文件。
make_txt.py文件
import os
import random
trainval_percent = 0.5
train_percent = 0.5
xmlfilepath = './Annotations/'
txtsavepath = './ImageSets/Main/'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
ftrainval = open(txtsavepath + '/trainval.txt', 'w')
ftest = open(txtsavepath + '/test.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fval = open(txtsavepath + '/val.txt', 'w')
for i in list:
name = total_xml[i][:-4] + '\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
(二)下载源码并修改
第一步:直接去GitHub上下载SSD源码:https://github.com/amdegroot/ssd.pytorch,我是直接dowmload到D盘。解压之后在打开SSD/data文件夹,在这之下新建一个VOCdevkit文件夹,然后把(一)中做好的VOC2007数据集copy到里面。
第二步:对源码进行修改。
用parcharm打开SSD项目。
修改一:data/init.py程序:将第三行的from .coco import COCODetection, COCOAnnotationTransform, COCO_CLASSES, COCO_ROOT, get_label_map
注释掉。因为我们不使用coco数据集,而是使用自己训练的。
修改二:data/config.py程序:第十五行的 num_classes
改成自己设定的类别数+1,所以我这里是6。
修改三:data/voc0712.py:第二十行的VOC_CLASSES =
改成自己的类别名,我的是bag,clock,phone,cup,mouse
;第二十三行的代码改成VOC_ROOT = './data/VOCdevkit/'
;第93行改为image_sets=[('2007', 'trainval')]
,因为我们只有2007数据集。
修改四:layers/modules/multibox_loss.py程序:在第97行的loss_c[pos] = 0
的全面加上一句loss_c = loss_c.view(num, -1)
,在第115行N = num_pos.data.sum()
的后面加上.double
,继续在下面添加loss_l = loss_l.double()
loss_c = loss_c.double()
这两句代码。
修改五:ssd.py程序:把这里面所有的num_classes
的数量都改为类别数+1,我的都是6。
修改六:train.py程序:搜索这里面的data[0]
,把这里面的[0]
都删掉,只保留data
。第84行的这两行#if args.dataset_root == COCO_ROOT: # parser.error('Must specify dataset if specifying dataset_root')
注释掉,因为我们不用COCO。
第165行的images, targets = next(batch_iterator)
改成
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
第198行的5000我改成500;迭代500次保存一次文件,大家可改可不改。第200行权重的名称自己随意,我改的是weights/ssd300_voc_
。
(三)开始训练
开始训练时需要一个预训练文件vgg16_reducedfc.pth,
百度云链接:https://pan.baidu.com/s/1FfqbkSe8spgxzYtCtPBAOQ
提取码:xg4c
大家可自行下载。
下载之后把他放在SSD项目下新建的weights文件夹下,然后就可以进行训练了。
注:训练中途遇到loss=nan的现象,我的解决方法是将train.py中,parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float,
中的 default=1e-3改为default=1e-4。
每500次保存一下文件。直到loss降低到1左右时即可。我训练了2000次,在权重weights得到如下所示文件:
我选了最后一次的文件进行测试,测试代码如下:
demo.py文件,放在demo文件夹里,需要自己修改的已标注
import os
import sys
import torch
from torch.autograd import Variable
import numpy as np
import cv2
from ssd import build_ssd
from data import VOC_CLASSES as labels
from matplotlib import pyplot as plt
module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
sys.path.append(module_path)
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net = build_ssd('test', 300, 6) # 第一处修改:类别+1
# 将预训练的权重加载到数据集上
net.load_weights('../weights/ssd300_voc_2000.pth') # 第二处修改:使用自己训练好的文件
# 加载多张图像
imgs = '../test_image/'# 第三处修改:改成你自己的文件夹
img_list = os.listdir(imgs)
for img in img_list:
# 对输入图像进行预处理
current_img = imgs + img
image = cv2.imread(current_img)
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
x = cv2.resize(image, (300, 300)).astype(np.float32)
x -= (104.0, 117.0, 123.0)
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
x = torch.from_numpy(x).permute(2, 0, 1)
# 把图片设为变量
xx = Variable(x.unsqueeze(0))
if torch.cuda.is_available():
xx = xx.cuda()
y = net(xx)
# 解析 查看结果
top_k = 10
plt.figure(figsize=(6, 6))
colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
currentAxis = plt.gca()
detections = y.data
scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)
for i in range(detections.size(1)):
j = 0
while detections[0, i, j, 0] >= 0.6:
score = detections[0, i, j, 0]
label_name = labels[i-1]
display_txt = '%s: %.2f'%(label_name, score)
print(display_txt)
pt = (detections[0,i,j,1:]*scale).cpu().numpy()
coords = (pt[0], pt[1]), pt[2]-pt[0]+1, pt[3]-pt[1]+1
color = colors[i]
currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor':color, 'alpha':0.5})
j += 1
plt.imshow(rgb_image)
plt.show()
(四)结果
本文主要参考博主陈纾的方法,十分感谢!
上一篇: VS2017搭建opencv开发环境
下一篇: weex-html5 组件进阶