test_net_state = torch.load(os.path.join(results_path))
test_net.load_state_dict(test_net_state)
if torch.cuda.is_available():
test_net.cuda()
test_net.eval()
# 3. Run Forward -------------------------------------------------------------------
with torch.no_grad():
pred_scores_tensor, pred_bbox_tensor = test_net.forward(
img_tensor.unsqueeze(0)) # N C H W
prior = CityScapeDataset([])
prior_bbox = prior.get_prior_bbox()
pred_scores_tensor = F.softmax(
pred_scores_tensor,
dim=2) # eval mode softmax was disabled in ssd_test
pred_bbox_tensor = loc2bbox(
pred_bbox_tensor,
CityScapeDataset([]).get_prior_bbox().unsqueeze(0))
pred_picked = nms_bbox(pred_bbox_tensor[0], pred_scores_tensor[0])
# 4. plot result
test_image = img_tensor.cpu().numpy().astype(
np.float32).transpose().copy() # H, W, C
test_image = ((test_image + 1) / 2)
for cls_dict in pred_picked:
# Test process
if WILL_TEST:
test_net = SSD(num_classes=3)
test_net.cuda()
test_net_state = torch.load('ssd_net.pth')
test_net.load_state_dict(test_net_state)
test_net.eval()
for test_batch_idx, (test_loc_targets, test_conf_targets,
test_image) in enumerate(test_data_loader):
# test_batch_idx, (test_loc_targets, test_conf_targets, test_image) = next(enumerate(test_data_loader))
test_image_permuted = test_image.permute(0, 3, 1, 2)
test_image_permuted = Variable(test_image_permuted.cuda().float())
test_conf_preds, test_loc_preds = test_net.forward(test_image_permuted)
# uset CityScapeDataset function get_prior_bbox
test_bbox_priors = test_set.get_prior_bbox().unsqueeze(0)
test_bbox_preds = loc2bbox(test_loc_preds.cpu(),
test_bbox_priors.cpu(),
center_var=0.1,
size_var=0.2)
sel_bbox_preds = nms_bbox(test_bbox_preds.squeeze().detach(),
test_conf_preds.squeeze().detach().cpu(),
overlap_threshold=0.5,
prob_threshold=0.5)
rects = []
texts = []
for key in sel_bbox_preds.keys():
for value in sel_bbox_preds[key]:
texts.append(key)
rects.append(value)
rects = center2corner(torch.tensor(rects)) * 300