【pytorch-ssd目标检测】测试自己创建的数据集
制作类似pascal voc格式的目标检测数据集:https://www.cnblogs.com/xiximayou/p/12546061.html
训练自己创建的数据集:https://www.cnblogs.com/xiximayou/p/12546556.html
验证自己创建的数据集:https://www.cnblogs.com/xiximayou/p/12550471.html
直接看修改后的text.py:
from __future__ import print_function import sys import os import argparse import torch import torch.nn as nn import torch.backends.cudnn as cudnn import torchvision.transforms as transforms from torch.autograd import Variable #from data import VOC_ROOT, VOC_CLASSES as labelmap from data import MASK_ROOT, MASK_CLASSES as labelmap from PIL import Image #from data import VOCAnnotationTransform, VOCDetection, BaseTransform, VOC_CLASSES from data import MASKAnnotationTransform, MASKDetection, BaseTransform, MASK_CLASSES import torch.utils.data as data from ssd import build_ssd parser = argparse.ArgumentParser(description='Single Shot MultiBox Detection') parser.add_argument('--trained_model', default='weights/ssd300_MASK_5000.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') parser.add_argument('--cuda', default=True, type=bool, help='Use cuda to train model') #parser.add_argument('--voc_root', default=VOC_ROOT, help='Location of VOC root directory') parser.add_argument('-f', default=None, type=str, help="Dummy arg so we can load in Jupyter Notebooks") args = parser.parse_args() if args.cuda and torch.cuda.is_available(): torch.set_default_tensor_type('torch.cuda.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) def test_net(save_folder, net, cuda, testset, transform, thresh): # dump predictions and assoc. ground truth to text file for now filename = save_folder+'test1.txt' num_images = len(testset) for i in range(num_images): print('Testing image {:d}/{:d}....'.format(i+1, num_images)) img = testset.pull_image(i) img_id, annotation = testset.pull_anno(i) x = torch.from_numpy(transform(img)[0]).permute(2, 0, 1) x = Variable(x.unsqueeze(0)) with open(filename, mode='a') as f: f.write('\nGROUND TRUTH FOR: '+img_id+'\n') for box in annotation: f.write('label: '+' || '.join(str(b) for b in box)+'\n') if cuda: x = x.cuda() y = net(x) # forward pass detections = y.data # scale each detection back up to the image scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]) pred_num = 0 for i in range(detections.size(1)): j = 0 while detections[0, i, j, 0] >= 0.6: if pred_num == 0: with open(filename, mode='a') as f: f.write('PREDICTIONS: '+'\n') score = detections[0, i, j, 0] label_name = labelmap[i-1] pt = (detections[0, i, j, 1:]*scale).cpu().numpy() coords = (pt[0], pt[1], pt[2], pt[3]) pred_num += 1 with open(filename, mode='a') as f: f.write(str(pred_num)+' label: '+label_name+' score: ' + str(score) + ' '+' || '.join(str(c) for c in coords) + '\n') j += 1 def test_voc(): # load net num_classes = len(MASK_CLASSES) + 1 # +1 background net = build_ssd('test', 300, num_classes) # initialize SSD net.load_state_dict(torch.load(args.trained_model)) net.eval() print('Finished loading model!') # load data mask_root="/content/drive/My Drive/pytorch_ssd" testset = MASKDetection(mask_root, "test", None, MASKAnnotationTransform()) if args.cuda: net = net.cuda() cudnn.benchmark = True # evaluation test_net(args.save_folder, net, args.cuda, testset, BaseTransform(net.size, (104, 117, 123)), thresh=args.visual_threshold) if __name__ == '__main__': test_voc()
开始执行:
!python test.py --trained_model weights/ssd300_MASK_5000.pth
运行结果:
Finished loading model! Testing image 1/80.... /pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function) Testing image 2/80.... 。。。。。。 /pytorch/torch/csrc/autograd/python_function.cpp:648: UserWarning: Legacy autograd function object was called twice. You will probably get incorrect gradients from this computation, as the saved tensors from the second invocation will clobber the saved tensors from the first invocation. Please consider rewriting your autograd function in the modern style; for information on the new format, please see: https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd Testing image 80/80....
看下生成了的文件:
看下test1.py中是什么:
GROUND TRUTH FOR: test_00000007 label: 46.0 || 0.0 || 139.0 || 128.0 || 0 PREDICTIONS: 1 label: mask score: tensor(0.9097) 31.465145 || 5.5611525 || 149.25903 || 86.10434 GROUND TRUTH FOR: test_00000010 label: 24.0 || 9.0 || 113.0 || 133.0 || 0 PREDICTIONS: 1 label: mask score: tensor(0.8791) 21.426735 || 17.9471 || 112.9484 || 122.676765 GROUND TRUTH FOR: test_00000015 label: 407.0 || 37.0 || 486.0 || 143.0 || 0 PREDICTIONS: 1 label: mask score: tensor(0.8441) 403.54123 || 42.476467 || 487.46075 || 146.36295 GROUND TRUTH FOR: test_00000016 label: 156.0 || 135.0 || 277.0 || 265.0 || 0 PREDICTIONS: 1 label: mask score: tensor(0.9541) 159.74387 || 109.33117 || 284.67053 || 264.61325
。。。。。。
每一张图片的坐标、置信度。