itr += 1
imgs = imgs.permute(0, 3, 1, 2).contiguous()
if use_gpu:
imgs = imgs.cuda()
imgs = Variable(imgs)
conf, loc = net.forward(imgs)
conf = conf[0, ...]
loc = loc[0, ...].cpu()
prior = test_dataset.get_prior_bbox()
prior = torch.unsqueeze(prior, 0)
# prior = prior.cuda()
real_bounding_box = loc2bbox(loc, prior, center_var=0.1, size_var=0.2)
real_bounding_box = torch.squeeze(real_bounding_box, 0)
class_list, sel_box = nms_bbox(real_bounding_box,
conf,
overlap_threshold=0.5,
prob_threshold=0.6)
# img = Image.open(os.path.join(img_dir, 'bad-honnef', 'bad-honnef_000000_000000_leftImg8bit.png'))
img = imgs[0].permute(1, 2, 0).contiguous()
true_loc = loc_targets[0, conf_targets[0, ...].nonzero(), :].squeeze()
img = Image.fromarray(np.uint8(img * 128 + 127))
# draw = ImageDraw.Draw(img)
sel_box = np.array(sel_box)
# loc_targets = torch.squeeze(loc2bbox(loc_targets, prior)).numpy()
fig, ax = plt.subplots(1)
ax.imshow(img)
for idx in range(len(sel_box)):
cx, cy, w, h = sel_box[idx] * 300
import numpy as np import torch.nn import torch.optim as optim import cityscape_dataset as csd from torch.utils.data import Dataset from torch.autograd import Variable from bbox_helper import generate_prior_bboxes, match_priors,nms_bbox,loc2bbox from data_loader import get_list from ssd_net import SSD from bbox_loss import MultiboxLoss prior_box = train_dataset.get_prior_boxes() boxes = loc2bbox(loc_preds, prior_box, center_var=0.1, size_var=0.2) select_box = nms_bbox(boxes, conf_preds, overlap_threshold=0.5, prob_threshold=0.6)
bbox = bbox_helper.loc2bbox(loc_preds, prior_bboxes)
bbox = bbox[0].detach()
bbox_corner = bbox_helper.center2corner(bbox)
print(bbox_corner)
print(conf_preds)
print(bbox_corner.shape)
bbox_corner = bbox_corner
conf_preds = conf_preds[0].detach()
# idx = conf_preds[:, 2] > 0.6
# bbox_corner = bbox_corner[idx]
# bbox = bbox[idx]
# print(bbox_corner)
bbox_nms = bbox_helper.nms_bbox(bbox_corner, conf_preds)
print(bbox_nms[0])
bbox_nms = torch.Tensor(bbox_nms[0])
print(bbox_nms.shape)
bbox_nms_cen = bbox_helper.corner2center(bbox_nms)
test_img = test_img.detach()
channels = test_img.shape[1]
h, w = test_img.shape[2], test_img.shape[3]
img_r = test_img.reshape(h, w, channels)
img_n = (img_r + 1) / 2
fig, ax = plt.subplots(1)
ax.imshow(img_n)
if USE_GPU:
test_net_state = torch.load(path_to_trained_model)
else:
test_net_state = torch.load(path_to_trained_model, map_location='cpu')
test_net = SSD(num_classes=3)
test_net.load_state_dict(test_net_state)
test_net.eval()
test_image_permuted = img_tensor.permute(0, 3, 1, 2)
test_image_permuted = Variable(test_image_permuted.float())
test_conf_preds, test_loc_preds = test_net.forward(test_image_permuted)
test_bbox_priors = prior_bboxes.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 = []
classes = []
for key in sel_bbox_preds.keys():
for value in sel_bbox_preds[key]:
classes.append(key)
rects.append(value)
rects = center2corner(torch.tensor(rects))*300
if len(rects) != 0:
if rects.dim() == 1:
rects = rects.unsqueeze(0)
classes = torch.tensor(classes)
drawRectsWithImgPLT(img, rects, classes)
def test_net(test_dataset, class_labels, results_path):
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
# Load the save model and deploy
test_net = SSD(len(class_labels))
test_net_state = torch.load(os.path.join(results_path))
test_net.load_state_dict(test_net_state)
test_net.cuda()
test_net.eval()
# accuracy
count_matched = 0
count_gt = 0
for test_item_idx in range(0, len(test_dataset)):
# test_item_idx = random.choice(range(0, len(test_dataset)))
test_image_tensor, test_label_tensor, test_bbox_tensor, prior_bbox = test_dataset[
test_item_idx]
# run Forward
with torch.no_grad():
pred_scores_tensor, pred_bbox_tensor = test_net.forward(
test_image_tensor.unsqueeze(0).cuda()) # N C H W
# scores -> Prob # because I deleted F.softmax~ at the ssd_net for net.eval
pred_scores_tensor = F.softmax(pred_scores_tensor, dim=2)
# bbox loc -> bbox (center)
pred_bbox_tensor = loc2bbox(pred_bbox_tensor, prior_bbox.unsqueeze(0))
# NMS : return tensor dictionary (bbo
pred_picked = nms_bbox(
pred_bbox_tensor[0],
pred_scores_tensor[0]) # not tensor, corner form
# Show the result
test_image = test_image_tensor.cpu().numpy().astype(
np.float32).transpose().copy() # H, W, C
test_image = ((test_image + 1) / 2)
gt_label = test_label_tensor.cpu().numpy().astype(np.uint8).copy()
gt_bbox_tensor = torch.cat([
test_bbox_tensor[..., :2] - test_bbox_tensor[..., 2:] / 2,
test_bbox_tensor[..., :2] + test_bbox_tensor[..., 2:] / 2
],
dim=-1)
gt_bbox = gt_bbox_tensor.detach().cpu().numpy().astype(
np.float32).reshape((-1, 4)).copy() * 300
gt_idx = gt_label > 0
# Calculate accuracy
pred_scores = pred_scores_tensor.detach().cpu().numpy().astype(
np.float32).copy()
pred_label = pred_scores[0].argmax(axis=1)
n_matched = 0
for gt, pr in zip(gt_label, pred_label):
if gt > 0 and gt == pr:
n_matched += 1
acc_per_image = 100 * n_matched / gt_idx.sum()
count_matched += n_matched
count_gt += gt_idx.sum()
# Show the results
gt_bbox = gt_bbox[gt_idx]
gt_label = gt_label[gt_idx]
if False:
for idx in range(gt_bbox.shape[0]):
cv2.rectangle(test_image, (gt_bbox[idx][0], gt_bbox[idx][1]),
(gt_bbox[idx][2], gt_bbox[idx][3]), (255, 0, 0),
1)
cv2.putText(test_image, str(gt_label[idx]),
(gt_bbox[idx][0], gt_bbox[idx][1]),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 0, 0), 1,
cv2.LINE_AA)
#--------------------
# cv2.rectangle(test_image, (pred_bbox[idx][0], pred_bbox[idx][1]), (pred_bbox[idx][2], pred_bbox[idx][3]),
# (0, 255, 0), 1)
#-----------------------
for cls_dict in pred_picked:
for p_score, p_bbox in zip(cls_dict['picked_scores'],
cls_dict['picked_bboxes']):
p_lbl = '%d | %.2f' % (cls_dict['class'], p_score)
p_bbox = p_bbox * 300
print(p_bbox, p_lbl)
cv2.rectangle(test_image, (p_bbox[0], p_bbox[1]),
(p_bbox[2], p_bbox[3]), (0, 0, 255), 2)
cv2.putText(test_image, p_lbl, (p_bbox[0], p_bbox[1]),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1,
cv2.LINE_AA)
plt.imshow(test_image)
plt.suptitle(class_labels)
plt.title('Temp Accuracy: {} %'.format(acc_per_image))
plt.show()
acc = 100 * count_matched / count_gt
print('Classification acc: ', '%')
# 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:
for p_score, p_bbox in zip(cls_dict['picked_scores'],
cls_dict['picked_bboxes']):
p_lbl = '%s | %.2f' % (class_labels[cls_dict['class']], p_score)
p_bbox = p_bbox * 300
print(p_bbox, p_lbl)
cv2.rectangle(test_image, (p_bbox[0], p_bbox[1]),
(p_bbox[2], p_bbox[3]), (0, 0, 255), 2)
cv2.putText(test_image, p_lbl, (p_bbox[0], p_bbox[1] - 3),