def main():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.multiprocessing.set_start_method('spawn')
data_list = load_data(cityscape_img_dir, cityscape_label_dir)
random.shuffle(data_list)
num_total_items = len(data_list)
net = SSD(5)
# Training set, ratio: 80%
num_train_sets = 0.8 * num_total_items
train_set_list = data_list[:int(num_train_sets)]
validation_set_list = data_list[int(num_train_sets):]
# Create dataloaders for training and validation
train_dataset = CityScapeDataset(train_set_list)
train_data_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=8,
shuffle=True,
num_workers=0)
print('Total training items',
len(train_dataset), ', Total training mini-batches in one epoch:',
len(train_data_loader))
validation_dataset = CityScapeDataset(validation_set_list)
validation_data_loader = torch.utils.data.DataLoader(validation_dataset,
batch_size=8,
shuffle=True,
num_workers=0)
print('Total validation items:', len(validation_dataset))
if Tuning:
net_state = torch.load(os.path.join(pth_path, 'ssd_net.pth'))
print('Loading trained model: ', os.path.join(pth_path, 'ssd_net.pth'))
net.load_state_dict(net_state)
train(net, train_data_loader, validation_data_loader)
test_list = get_list(img_dir, label_dir)
# test_list = test_list[0:-20]
test_dataset = csd.CityScapeDataset(test_list, train=False, show=False)
test_data_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=16,
shuffle=False,
num_workers=0)
print('test items:', len(test_dataset))
file_name = 'SSD'
test_net_state = torch.load(os.path.join('.', file_name + '.pth'))
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()
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())
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: ', '%')
test_list = test_list[0:10]
test_dataset = csd.CityScapeDataset(test_list, train=False, show=False)
test_data_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=0)
print('test items:', len(test_dataset))
net = SSD(3)
if use_gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net.cuda()
net_state = torch.load(os.path.join(save_dir, file_name + '.pth'))
net.load_state_dict(net_state)
net.eval()
for test_batch_idx, (loc_targets, conf_targets,
imgs) in enumerate(test_data_loader):
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)
real_bounding_box = loc2bbox(loc, prior)