欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测

程序员文章站 2024-03-16 10:40:34
...

Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测

(0)环境

Windows10
PyTorch1.4
pycharm

(一)准备数据集

首先新建了VOC2007文件夹,然后再在里面新建三个文件夹,如下图。Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
其中,JPEGImages用来存放照片,Annotations存放xml文件。
第一步:对图片进行标注。我这里选了五种类别(书包,钟表,手机,水杯,鼠标)各十张图片,总共五十张照片(由于显卡不行,只选了这么多),标注工具为labelImg,教程网上很多,这里不再赘述。标注得到的xml文件存放在Annotations中。
Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
第二步:用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()

Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测

(二)下载源码并修改

第一步:直接去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得到如下所示文件:
Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
我选了最后一次的文件进行测试,测试代码如下:
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()

(四)结果

Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
Windows下自己训练数据集,在pytorch上实现基于SSD的物体检测
本文主要参考博主陈纾的方法,十分感谢!