手把手教物体检测——SSD（新手必看）
-
下载SSD代码：
https://github.com/amdegroot/ssd.pytorch
- 将下载的代码解压后在data文件夹下新建VOCdevkit文件夹，然后将VOC2007数据集复制到该文件夹下面。
- 下载权重文件放在weights文件夹下面。下载地址：
https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
- 修改config.py代码
-
# SSD300 CONFIGS
-
-
voc = {
-
-
'num_classes': 3,//将类别改为：类别+1（背景）
-
-
'lr_steps': (80000, 100000, 120000),
-
-
'max_iter': 120000,//迭代次数
-
-
'feature_maps': [38, 19, 10, 5, 3, 1],
-
-
'min_dim': 300,
-
-
'steps': [8, 16, 32, 64, 100, 300],
-
-
'min_sizes': [30, 60, 111, 162, 213, 264],
-
-
'max_sizes': [60, 111, 162, 213, 264, 315],
-
-
'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
-
-
'variance': [0.1, 0.2],
-
-
'clip': True,
-
-
'name': 'VOC',
-
-
}
- 修改VOC0712.py代码。
-
'''
-
-
VOC_CLASSES = ( # always index 0
-
-
'aeroplane', 'bicycle', 'bird', 'boat',
-
-
'bottle', 'bus', 'car', 'cat', 'chair',
-
-
'cow', 'diningtable', 'dog', 'horse',
-
-
'motorbike', 'person', 'pottedplant',
-
-
'sheep', 'sofa', 'train', 'tvmonitor')
-
-
'''
-
-
VOC_CLASSES = ( # always index 0
-
-
'aircraft', 'oiltank')//修改为自己数据集的类别。
-
-
# note: if you used our download scripts, this should be right
-
-
#VOC_ROOT = osp.join("", "data/VOCdevkit/")
-
-
VOC_ROOT = "data/VOCdevkit/" #修改为Win10路径
将
-
def __init__(self, root,
-
-
image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
-
-
transform=None, target_transform=VOCAnnotationTransform(),
-
-
dataset_name='VOC0712'):
修改为：
-
def __init__(self, root,
-
-
image_sets=[('2007', 'trainval')],
-
-
transform=None, target_transform=VOCAnnotationTransform(),
-
-
dataset_name='VOC2007'):
- 修改coco.py
COCO_ROOT ='data/'
- 修改ssd.py
将32行的self.cfg = (coco, voc)[num_classes == 21]
修改为self.cfg = (coco, voc)[num_classes == 3]#3是我数据集的类别+1
- 修改multibox_loss.py
在97行的loss_c[pos] = 0 # filter out pos boxes for now
位置之上加入loss_c = loss_c.view(pos.size()[0], pos.size()[1])
如果没有这句话会引起张量不匹配的问题
如：
IndexError: The shape of the mask [4, 8732] at index 0 does not match the shape of the indexed tensor [34928, 1] at index 0
-
修改全局配置参数：
parser.add_argument('--dataset', default='VOC', choices=['VOC', 'COCO'],
type=str, help='VOC or COCO')#改为VOC
parser.add_argument('--batch_size', default=4, type=int,
help='Batch size for training')#Batch Size 按照显存的大小设置合适的值。
parser.add_argument('--start_iter', default=0, type=int,
help='Resume training at this iter')#迭代开始
parser.add_argument('--num_workers', default=4, type=int,
help='Number of workers used in dataloading')
parser.add_argument('--cuda', default=True, type=str2bool,
help='Use CUDA to train model')
parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float,
help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float,
help='Momentum value for optim')
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('--visdom', default=False, type=str2bool,
help='Use visdom for loss visualization')
parser.add_argument('--save_folder', default='weights/',
help='Directory for saving checkpoint models')#保存权重的位置。
args = parser.parse_args()将：修改为：-
loc_loss += loss_l.item()
-
-
conf_loss += loss_c.item()
-
-
if iteration % 10 == 0:
-
-
print('timer: %.4f sec.' % (t1 - t0))
-
-
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ')
不修改会引起错误：.IndexError: invalid index of a 0-dim tensor. Use tensor.item() to convert a 0-dim tensor to a Python number将165行的images, targets = next(batch_iterator)更改成：try:images, targets = next(batch_iterator)except StopIteration:batch_iterator = iter(data_loader)images, targets = next(batch_iterator)如果不改会引起跳出迭代的问题：Traceback (most recent call last):File "E:/ssd.pytorch-master/ssd.pytorch-master/train.py", line 246, in <module>train()File "E:/ssd.pytorch-master/ssd.pytorch-master/train.py", line 159, in trainimages, targets = next(batch_iterator)File "D:\Users\WH\anaconda3\lib\site-packages\torch\utils\data\dataloader.py", line 345, in __next__data = self._next_data()File "D:\Users\WH\anaconda3\lib\site-packages\torch\utils\data\dataloader.py", line 831, in _next_dataraise StopIterationStopIteration如果遇到loss为nan的现象，就需要降低出示学习率。将10e-3改为10e-4将以上问题解决后就可以开始训练了。 -
- 测试
需要修改test.py的全局参数。
parser.add_argument('--trained_model', default='weights/ssd300_COCO_60000.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('--visual_threshold', default=0.6, type=float,
help='Final confidence threshold')#检测的最低置信度。
点击测试生成测试结果：
- 打印测试结果
import os
import sys
module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
sys.path.append(module_path)
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import numpy as np
import cv2
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
from ssd import build_ssd
net = build_ssd('test', 300, 3) # initialize SSD
net.load_state_dict(torch.load('weights/ssd300_COCO_60000.pth'))#加载模型
net.eval()
image = cv2.imread('data/VOCdevkit/aircraft_27.jpg', cv2.IMREAD_COLOR) # 加载图片
from matplotlib import pyplot as plt
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)#读取图片
# View the sampled input image before transform
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()
plt.imshow(x)
x = torch.from_numpy(x).permute(2, 0, 1)
xx = Variable(x.unsqueeze(0)) # wrap tensor in Variable
if torch.cuda.is_available():
xx = xx.cuda()
y = net(xx)
from data import VOC_CLASSES as labels
top_k=10#选择前10个结果
plt.figure(figsize=(10,10))
colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
plt.imshow(rgb_image) # plot the image for matplotlib
currentAxis = plt.gca()
detections = y.data
# scale each detection back up to the image
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.35:
score = detections[0,i,j,0]
label_name = labels[i-1]
display_txt = '%s: %.2f'%(label_name, score)
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.show()
文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。
原文链接：wanghao.blog.csdn.net/article/details/105800157
- 点赞
- 收藏
- 关注作者
评论（0）