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():