net.train()
imgs = imgs.permute(
0, 3, 1,
2).contiguous() # [batch_size, W, H, CH] -> [batch_size, CH, W, H]
if use_gpu:
imgs = imgs.cuda()
loc_targets = loc_targets.cuda()
conf_targets = conf_targets.cuda()
imgs = Variable(imgs)
loc_targets = Variable(loc_targets)
conf_targets = Variable(conf_targets)
optimizer.zero_grad()
conf_preds, loc_preds = net.forward(imgs)
conf_loss, loc_huber_loss, loss = criterion.forward(
conf_preds, loc_preds, conf_targets, loc_targets)
loss.backward()
optimizer.step()
train_losses.append(
(itr, conf_loss.item(), loc_huber_loss.item(), loss.item()))
# if train_batch_idx % 2 == 0:
print('Epoch: %d Itr: %d Conf_Loss: %f Loc_Loss: %f Loss: %f' %
(epoch_idx, itr, conf_loss.item(), loc_huber_loss.item(),
loss.item()))
net_state = net.state_dict()
file_name = 'SSD'
net = SSD(3)
if use_gpu:
net = net.cuda()
net.load_state_dict(test_net_state)
itr = 0
net.eval()
for test_batch_idx, (loc_targets, conf_targets,
imgs) in enumerate(test_data_loader):
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()
prior_bboxes = generate_prior_bboxes(prior_layer_cfg = prior_layer_cfg)
if WILL_TEST:
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:
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: ', '%')
def main():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
prior_layer_cfg = [{
'layer_name': 'Conv5',
'feature_dim_hw': (19, 19),
'bbox_size': (60, 60),
'aspect_ratio': (1.0, 1 / 2, 1 / 3, 2.0, 3.0, '1t')
}, {
'layer_name': 'Conv11',
'feature_dim_hw': (10, 10),
'bbox_size': (105, 105),
'aspect_ratio': (1.0, 1 / 2, 1 / 3, 2.0, 3.0, '1t')
}, {
'layer_name': 'Conv14_2',
'feature_dim_hw': (5, 5),
'bbox_size': (150, 150),
'aspect_ratio': (1.0, 1 / 2, 1 / 3, 2.0, 3.0, '1t')
}, {
'layer_name': 'Conv15_2',
'feature_dim_hw': (3, 3),
'bbox_size': (195, 195),
'aspect_ratio': (1.0, 1 / 2, 1 / 3, 2.0, 3.0, '1t')
}, {
'layer_name': 'Conv16_2',
'feature_dim_hw': (2, 2),
'bbox_size': (240, 240),
'aspect_ratio': (1.0, 1 / 2, 1 / 3, 2.0, 3.0, '1t')
}, {
'layer_name': 'Conv17_2',
'feature_dim_hw': (1, 1),
'bbox_size': (285, 285),
'aspect_ratio': (1.0, 1 / 2, 1 / 3, 2.0, 3.0, '1t')
}]
prior_bboxes = generate_prior_bboxes(prior_layer_cfg)
# loading the test image
img_file_path = sys.argv[1]
# img_file_path = 'image.png'
img = Image.open(img_file_path)
img = img.resize((300, 300))
plot_img = img.copy()
img_array = np.asarray(img)[:, :, :3]
mean = np.asarray((127, 127, 127))
std = 128.0
img_array = (img_array - mean) / std
h, w, c = img_array.shape[0], img_array.shape[1], img_array.shape[2]
img_tensor = torch.Tensor(img_array)
test_input = img_tensor.view(1, c, h, w)
# # loading test input to run test on
# test_data_loader = torch.utils.data.DataLoader(test_input,
# batch_size=1,
# shuffle=True,
# num_workers=0)
# idx, (img) = next(enumerate(test_data_loader))
# # Setting model to evaluate mode
net = SSD(2)
test_net_state = torch.load('ssd_net.pth')
net.load_state_dict(test_net_state)
# net.eval()
net.cuda()
# Forward
test_input = Variable(test_input.cuda())
test_cof, test_loc = net.forward(test_input)
test_loc = test_loc.detach()
test_loc_clone = test_loc.clone()
# normalizing the loss to add up to 1 (for probability)
test_cof_score = F.softmax(test_cof[0], dim=1)
# print(test_cof_score.shape)
# print(test_cof_score)
# running NMS
sel_idx = nms_bbox1(test_loc_clone[0],
prior_bboxes,
test_cof_score.detach(),
overlap_threshold=0.5,
prob_threshold=0.24)
test_loc = loc2bbox(test_loc[0], prior_bboxes)
test_loc = center2corner(test_loc)
sel_bboxes = test_loc[sel_idx]
# plotting the output
plot_output(plot_img, sel_bboxes.cpu().detach().numpy())
# 2. Load the saved model and test ------------------------------------------------
class_labels = list(dataset_label_group.keys())
test_net = SSD(len(class_labels))
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
valid_losses = []
max_epochs = 50
itr = 0
# Train process
for epoch_idx in range(0, max_epochs):
for train_batch_idx, (train_loc_targets, train_conf_targets,
train_image) in enumerate(train_data_loader):
itr += 1
net.train()
train_loc_targets = Variable(train_loc_targets.cuda().float())
train_conf_targets = Variable(train_conf_targets.cuda().long())
train_image = train_image.permute(0, 3, 1, 2)
train_image = Variable(train_image.cuda().float())
optimizer.zero_grad()
train_conf_preds, train_loc_preds = net.forward(train_image)
train_conf_loss, train_loc_loss = criterion.forward(
train_conf_preds, train_loc_preds, train_conf_targets,
train_loc_targets)
train_total_loss = train_conf_loss + train_loc_loss
train_total_loss.backward()
optimizer.step()
train_losses.append((itr, train_total_loss))
print(
'Train Epoch: %d Itr: %d Conf Loss: %f Loc Loss: %f Total Loss: %f'
% (epoch_idx, itr, train_conf_loss.item(),
train_loc_loss.item(), train_total_loss.item()))
# Validation process:
if train_batch_idx % 5 == 0:
net.eval()