det_file = os.path.join(save_folder, 'detections.pkl') # load data if args.dataset == 'VOC': testset = VOCDetection(VOCroot, [('2007', 'test')], None, AnnotationTransform()) elif args.dataset == 'COCO': testset = COCODetection(COCOroot, [('2014', 'minival')], None) #COCOroot, [('2015', 'test-dev')], None) else: print('Only VOC and COCO dataset are supported now!') if args.retest: f = open(det_file, 'rb') all_boxes = pickle.load(f) print('Evaluating detections') testset.evaluate_detections(all_boxes, save_folder) else: # load net img_dim = (300, 512)[args.size == '512'] num_classes = (21, 81)[args.dataset == 'COCO'] net = build_net('test', img_dim, num_classes) # initialize detector state_dict = torch.load(args.trained_model, map_location='cuda:{}'.format(gpu_id)) # create new OrderedDict that does not contain `module.` from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7] if head == 'module.': name = k[7:] # remove `module.`
def do_test(args, model, detector, max_per_image=200, thresh=0.01): if args.dataset == 'VOC': dataset = VOCDetection( args, VOCroot, [('2007', 'test')], None, AnnotationTransform(0 if args.setting == 'transfer' else args.split), True) elif args.dataset == 'COCO': dataset = COCODetection(COCOroot, [('2014', 'split_nonvoc_minival')], None) else: raise ValueError(f"Unknown dataset: {args.dataset}") num_images = len(dataset) all_boxes = [[[] for _ in range(num_images)] for _ in range(num_classes)] transform = BaseTransform(model.size, rgb_means, (2, 0, 1)) _t = {'im_detect': Timer(), 'misc': Timer()} det_file = os.path.join(args.save_folder, 'detections.pkl') if args.retest: f = open(det_file, 'rb') all_boxes = pickle.load(f) logger.info('Evaluating detections') dataset.evaluate_detections(all_boxes, args.save_folder) return for i in range(num_images): img = dataset.pull_image(i) scale = torch.Tensor( [img.shape[1], img.shape[0], img.shape[1], img.shape[0]]).to(model.device) with torch.no_grad(): x = transform(img).unsqueeze(0) _t['im_detect'].tic() pred = model(x) # forward pass boxes, scores = detector.forward(pred, priors) detect_time = _t['im_detect'].toc() boxes = boxes[0] # percent and point form detection boxes scores = scores[0] # [1, num_priors, num_classes] boxes *= scale # scale each detection back up to the image boxes = boxes.cpu().numpy() scores = scores.cpu().numpy() _t['misc'].tic() for j in range(1, num_classes): inds = np.where(scores[:, j] > thresh)[0] if len(inds) == 0: all_boxes[j][i] = np.empty([0, 5], dtype=np.float32) continue 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, :] all_boxes[j][i] = c_dets if max_per_image > 0: image_scores = np.hstack( [all_boxes[j][i][:, -1] for j in range(1, num_classes)]) if len(image_scores) > max_per_image: image_thresh = np.sort(image_scores)[-max_per_image] for j in range(1, num_classes): keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0] all_boxes[j][i] = all_boxes[j][i][keep, :] nms_time = _t['misc'].toc() if i % 20 == 0: logger.info('im_detect: {:d}/{:d} {:.3f}s {:.3f}s'.format( i + 1, num_images, detect_time, nms_time)) _t['im_detect'].clear() _t['misc'].clear() with open(det_file, 'wb') as f: pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL) logger.info('Evaluating detections') dataset.evaluate_detections(all_boxes, args.save_folder)