def test_net(save_folder, net, detector, cuda, testset, transform, max_per_image=300, thresh=0.005): if not os.path.exists(save_folder): os.mkdir(save_folder) # dump predictions and assoc. ground truth to text file for now num_images = len(testset) print_info('=> Total {} images to test.'.format(num_images), ['yellow', 'bold']) num_classes = 2 all_boxes = [[[] for _ in range(num_images)] for _ in range(num_classes)] _t = {'im_detect': Timer(), 'misc': Timer()} det_file = os.path.join(save_folder, 'detections.pkl') print_info('Begin to evaluate', ['yellow', 'bold']) for i in tqdm(range(num_images)): img = testset.pull_image(i) scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]) with torch.no_grad(): x = transform(img).unsqueeze(0) if cuda: x = x.cuda() scale = scale.cuda() _t['im_detect'].tic() out = net(x) # forward pass boxes, scores = detector.forward(out,priors) detect_time = _t['im_detect'].toc() boxes = boxes[0] scores=scores[0] boxes *= scale boxes = boxes.cpu().numpy() scores = scores.cpu().numpy() # scale each detection back up to the image _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: print('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) print('Evaluating detections') testset.evaluate_detections(all_boxes, save_folder)
if __name__ == '__main__': # load net img_dim = 300 num_classes = 2 rgb_means = (104, 117, 123) start_epoch = 20 trained_model_list = os.listdir(args.weights_path) trained_model_list.sort() net = build_net('test', img_dim, num_classes) # initialize detector for trained_model in trained_model_list: start_epoch += 10 if start_epoch < 1 + 10: continue trained_model = os.path.join(args.weights_path, trained_model) print_info(trained_model, ['yellow', 'bold']) state_dict = torch.load(trained_model) # 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.` else: name = k new_state_dict[name] = v net.load_state_dict(new_state_dict) net.eval() print('Finished loading model!')
# parser.add_argument('-m', '--trained_model', default='weights/RFB_vgg_VOC_epoches_150.pth', # type=str, help='Trained state_dict file path to open') parser.add_argument('-m', '--trained_model', default='weights/7690.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('--cuda', default=True, type=bool, help='Use cuda to train model') parser.add_argument('--cpu', default=False, type=bool, help='Use cpu nms') parser.add_argument('--retest', default=False, type=bool, help='test cache results') args = parser.parse_args() print_info('----------------------------------------------------------------------\n' '| RFBDet Evaluation Program |\n' '----------------------------------------------------------------------', ['yellow','bold']) if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) cfg = VOC_300 if args.version == 'RFB_vgg': from models.RFB_Net_vgg import build_net elif args.version == 'RFB_E_vgg': from models.RFB_Net_E_vgg import build_net priorbox = PriorBox(cfg) with torch.no_grad():