def test():
import os
im_file = 'demo/004545.jpg'
# im_file = 'data/VOCdevkit2007/VOC2007/JPEGImages/009036.jpg'
# im_file = '/media/longc/Data/data/2DMOT2015/test/ETH-Crossing/img1/000100.jpg'
image = cv2.imread(im_file)
model_file = 'models/saved_model3/faster_rcnn_200000.pth'
# model_file = '/media/longc/Data/models/faster_rcnn_pytorch3/faster_rcnn_100000.h5'
# model_file = '/media/longc/Data/models/faster_rcnn_pytorch2/faster_rcnn_2000.h5'
detector = FasterRCNN()
detector.load_state_dict(torch.load(model_file))
detector.cuda()
detector.eval()
print('load model successfully!')
# network.save_net(r'/media/longc/Data/models/VGGnet_fast_rcnn_iter_70000.h5', detector)
# print('save model succ')
t = Timer()
t.tic()
# image = np.zeros(shape=[600, 800, 3], dtype=np.uint8) + 255
dets, scores, classes = detector.detect(image, 0.7)
runtime = t.toc()
print('total spend: {}s'.format(runtime))
im2show = np.copy(image)
for i, det in enumerate(dets):
det = tuple(int(x) for x in det)
cv2.rectangle(im2show, det[0:2], det[2:4], (255, 205, 51), 2)
cv2.putText(im2show, '%s: %.3f' % (classes[i], scores[i]), (det[0], det[1] + 15), cv2.FONT_HERSHEY_PLAIN,
1.0, (0, 0, 255), thickness=1)
cv2.imwrite(os.path.join('demo', 'out.jpg'), im2show)
cv2.imshow('demo', im2show)
cv2.waitKey(0)
def test():
args = parse_args()
# perpare data
print('load data')
if args.dataset == 'voc07test':
dataset_name = 'voc_2007_test'
elif args.dataset == 'voc12test':
dataset_name = 'voc_2012_test'
else:
raise NotImplementedError
cfg.TRAIN.USE_FLIPPED = False
imdb, roidb = combined_roidb(dataset_name)
test_dataset = RoiDataset(roidb)
test_dataloader = DataLoader(dataset=test_dataset,
batch_size=1,
shuffle=False)
test_data_iter = iter(test_dataloader)
# load model
model = FasterRCNN(backbone=args.backbone)
model_name = '0712_faster_rcnn101_epoch_{}.pth'.format(args.check_epoch)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
model_path = os.path.join(args.output_dir, model_name)
model.load_state_dict(torch.load(model_path)['model'])
if args.use_gpu:
model = model.cuda()
model.eval()
num_images = len(imdb.image_index)
det_file_path = os.path.join(args.output_dir, 'detections.pkl')
all_boxes = [[[] for _ in range(num_images)]
for _ in range(imdb.num_classes)]
empty_array = np.transpose(np.array([[], [], [], [], []]), (1, 0))
torch.set_grad_enabled(False)
for i in range(num_images):
start_time = time.time()
im_data, gt_boxes, im_info = next(test_data_iter)
if args.use_gpu:
im_data = im_data.cuda()
gt_boxes = gt_boxes.cuda()
im_info = im_info.cuda()
im_data_variable = Variable(im_data)
det_tic = time.time()
rois, faster_rcnn_cls_prob, faster_rcnn_reg, _, _, _, _, _ = model(
im_data_variable, gt_boxes, im_info)
scores = faster_rcnn_cls_prob.data
boxes = rois.data[:, 1:]
boxes_deltas = faster_rcnn_reg.data
if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:
boxes_deltas = boxes_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
+ torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
boxes_deltas = boxes_deltas.view(-1, 4 * imdb.num_classes)
pred_boxes = bbox_transform_inv_cls(boxes, boxes_deltas)
pred_boxes = clip_boxes_cls(pred_boxes, im_info[0])
pred_boxes /= im_info[0][2].item()
det_toc = time.time()
detect_time = det_tic - det_toc
nms_tic = time.time()
if args.vis:
im_show = Image.open(imdb.image_path_at(i))
for j in range(1, imdb.num_classes):
inds = torch.nonzero(scores[:, j] > args.thresh).view(-1)
if inds.numel() > 0:
cls_score = scores[:, j][inds]
_, order = torch.sort(cls_score, 0, True)
cls_boxes = pred_boxes[inds][:, j * 4:(j + 1) * 4]
cls_dets = torch.cat((cls_boxes, cls_score.unsqueeze(1)), 1)
cls_dets = cls_dets[order]
keep = nms(cls_dets, 0.3)
cls_dets = cls_dets[keep.view(-1).long()]
if args.vis:
cls_name_dets = np.repeat(j, cls_dets.size(0))
im_show = draw_detection_boxes(im_show,
cls_dets.cpu().numpy(),
cls_name_dets, imdb.classes,
0.5)
all_boxes[j][i] = cls_dets.cpu().numpy()
else:
all_boxes[j][i] = empty_array
if args.max_per_image > 0:
image_scores = np.hstack(
[all_boxes[j][i][:, -1] for j in range(1, imdb.num_classes)])
if len(image_scores) > args.max_per_image:
image_thresh = np.sort(image_scores)[-args.max_per_image]
for j in range(1, imdb.num_classes):
keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
all_boxes[j][i] = all_boxes[j][i][keep, :]
if args.vis:
plt.imshow(im_show)
plt.show()
nms_toc = time.time()
nms_time = nms_tic - nms_toc
sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s \r' \
.format(i + 1, num_images, detect_time, nms_time))
sys.stdout.flush()
with open(det_file_path, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
imdb.evaluate_detections(all_boxes, args.output_dir)
end_time = time.time()
print("test time: %0.4fs" % (end_time - start_time))
