def test_bbox2loc(self):
prior = torch.Tensor([[1, 2, 3, 4], [2, 0, 3, 4], [1, 5, 6, 7]])
gt = torch.Tensor([[2, 6, 7, 8], [5, 6, 7, 8], [2, 0, 4, 8]])
temp = bbox2loc(gt, prior)
temp = bbox2loc(gt, prior)
back = loc2bbox(temp, prior)
print(back)
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()
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()
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)
prior_bboxes = prior_bboxes.unsqueeze(0).cpu()
test_img = np.array(test_img)
test_img = np.subtract(test_img, [127, 127, 127])
test_img = np.divide(test_img, 128)
test_img = test_img.reshape(
(test_img.shape[2], test_img.shape[0], test_img.shape[1]))
test_img = torch.Tensor(test_img)
test_img = test_img.unsqueeze(0).cpu()
images = Variable(test_img) # Use Variable(*) to allow gradient flow\n",
conf_preds, loc_preds = net.forward(images) # Forward once\n",
print(conf_preds.shape)
print(loc_preds.shape)
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])
def test_random2(self):
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.set_printoptions(precision=10)
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')
}]
pp = generate_prior_bboxes(prior_layer_cfg)
# test_list = load_data('../Debugimage', '../Debuglabel')
test_list = load_data('../cityscapes_samples',
'../cityscapes_samples_labels')
#print(test_list)
test_dataset = CityScapeDataset(test_list)
test_data_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=True,
num_workers=0)
lfw_dataset_dir = '../'
test_net = ssd_net.SSD(3)
test_net_state = torch.load(
os.path.join(lfw_dataset_dir, 'ssd_net.pth'))
test_net.load_state_dict(test_net_state)
idx, (img, bbox, label) = next(enumerate(test_data_loader))
pred_cof, pred_loc = test_net.forward(img)
print(pred_loc.shape)
import torch.nn.functional as F
pred_loc = pred_loc.detach()
bbox_center = loc2bbox(pred_loc[0], pp)
pred_cof = F.softmax(pred_cof[0])
ind = np.where(pred_cof > 0.7)
# pred_cof = F.softmax(pred_cof[ind[0]])
bbox_center = bbox_center[ind[0]]
print(ind, pred_cof)
img = img[0].cpu().numpy()
img = img.reshape((300, 300, 3))
img = (img * 128 + np.asarray([[127, 127, 127]])) / 255
fig, ax = plt.subplots(1)
imageB_array = resize(img, (600, 1200), anti_aliasing=True)
ax.imshow(imageB_array, cmap='brg')
bbox_corner = center2corner(bbox_center)
for i in range(0, bbox_corner.shape[0]):
# print('i point', bbox_corner[i, 0]*600, bbox_corner[i, 1]*300,(bbox_corner[i, 2]-bbox_corner[i, 0])*600, (bbox_corner[i, 3]-bbox_corner[i, 1])*300)
rect = patches.Rectangle(
(bbox_corner[i, 0] * 1200, bbox_corner[i, 1] * 600),
(bbox_corner[i, 2] - bbox_corner[i, 0]) * 1200,
(bbox_corner[i, 3] - bbox_corner[i, 1]) * 600,
linewidth=2,
edgecolor='r',
facecolor='none') # Create a Rectangle patch
ax.add_patch(rect) # Add the patch to the Axes
plt.show()
def test_dataLoad(self):
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.set_printoptions(precision=10)
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')
}]
pp = generate_prior_bboxes(prior_layer_cfg)
#test_list = load_data('../Debugimage', '../Debuglabel')
test_list = load_data('../cityscapes_samples',
'../cityscapes_samples_labels')
print(test_list)
gt_bbox = np.asarray(test_list[0]['label'][1]) * [
300 / 2048, 300 / 1024, 300 / 2048, 300 / 1024
]
print('ground truth from file:', test_list[0]['label'][0])
test_dataset = CityScapeDataset(test_list)
test_data_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=True,
num_workers=0)
idx, (img, bbox, label) = next(enumerate(test_data_loader))
bbox = bbox[0]
label = label[0]
print(bbox.shape, label.shape)
print('matched label', label[np.where(label > 0)], np.where(label > 0),
label.shape)
print('first bbox from data_set:', bbox[0], label[0])
bbox_center = loc2bbox(bbox, pp)
bbox_corner = center2corner(bbox_center)
img = img[0].cpu().numpy()
img = img.reshape((300, 300, 3))
img = (img * 128 + np.asarray([[127, 127, 127]])) / 255
# for i in range(0, bbox.shape[0]):
# cv2.rectangle(img, (bbox[i,0], bbox[i,1]), (bbox[i,2], bbox[i,3]), (0, 255, 0), 3)
#cv2.imshow("img", img)
# Create figure and axes
fig, ax = plt.subplots(1)
imageB_array = resize(img, (300, 300), anti_aliasing=True)
ax.imshow(imageB_array, cmap='brg')
bbox_corner = bbox_corner.cpu().numpy()
bbox_corner = bbox_corner[np.where(label > 0)]
temp_lab = label[np.where(label > 0)]
print('matched bbox ======', bbox_corner)
pp = center2corner(pp)
pp = pp[np.where(label > 0)]
print('864 tensor: ', pp)
for i in range(0, bbox_corner.shape[0]):
if temp_lab[i] == 1:
# print('i point', bbox_corner[i, 0]*600, bbox_corner[i, 1]*300,(bbox_corner[i, 2]-bbox_corner[i, 0])*600, (bbox_corner[i, 3]-bbox_corner[i, 1])*300)
rect = patches.Rectangle(
(bbox_corner[i, 0] * 300, bbox_corner[i, 1] * 300),
(bbox_corner[i, 2] - bbox_corner[i, 0]) * 300,
(bbox_corner[i, 3] - bbox_corner[i, 1]) * 300,
linewidth=2,
edgecolor='r',
facecolor='none') # Create a Rectangle patch
ax.add_patch(rect) # Add the patch to the Axes
else:
rect = patches.Rectangle(
(bbox_corner[i, 0] * 1200, bbox_corner[i, 1] * 600),
(bbox_corner[i, 2] - bbox_corner[i, 0]) * 1200,
(bbox_corner[i, 3] - bbox_corner[i, 1]) * 600,
linewidth=2,
edgecolor='y',
facecolor='none') # Create a Rectangle patch
ax.add_patch(rect) # Add the patch to the Axes
for i in range(0, pp.shape[0]):
rect = patches.Rectangle(
(pp[i, 0] * 300, pp[i, 1] * 300), (pp[i, 2] - pp[i, 0]) * 300,
(pp[i, 3] - pp[i, 1]) * 300,
linewidth=1,
edgecolor='blue',
facecolor='none') # Create a Rectangle patch
ax.add_patch(rect) # Add the patch to the Axes
# for i in range(0, gt_bbox.shape[0]):
# rect = patches.Rectangle((gt_bbox[i][0], gt_bbox[i][1]),
# (gt_bbox[i][2] - gt_bbox[i][0]),
# (gt_bbox[i][3] - gt_bbox[i][1]), linewidth=1, edgecolor='g',
# facecolor='none') # Create a Rectangle patch
# ax.add_patch(rect) # Add the patch to the Axes
plt.show()
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:
rects = rects.unsqueeze(0)
classes = torch.tensor(classes)
shuffle=True,
num_workers=4)
for batch_idx, (img, bbox, label, priors) in enumerate(test_loader):
# print(batch_idx)
inputs = img.float().to(device)
pred_conf, pred_locs = test_model.forward(inputs)
# print(pred_conf)
# softmax = torch.nn.LogSoftmax(dim = 2)
# pred_conf = softmax(pred_conf)
for i in range(pred_locs.shape[0]):
single_pred_locs = pred_locs[i, :, :]
single_pred_conf = pred_conf[i, :, :]
# print('pred_locs:', single_pred_locs.shape)
# print('pred_conf:', single_pred_conf.shape)
single_pred_locs = loc2bbox(single_pred_locs,
priors[i:i + 1, :, :].cuda())[0, :, :]
# print(single_pred_locs.shape)
# for label_id in range(1,3):
# print(label_id)
img[i, :, :, :] = img[i, :, :, :] * 128.0 + 127.0
npimg = img[i, :, :, :].numpy().reshape(3, 300, 300)
npimg = npimg.astype(dtype='int')
fig, ax = plt.subplots(1)
ax.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')
for label_id in range(1, 3):
keep = nms(single_pred_locs, single_pred_conf[:, label_id])
# keep = single_pred_locs
# print('keep:', keep.shape)
# print(keepbase_output_layer_indices[0][0])
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())
def train_net(train_loader,
valid_loader,
class_labels,
lab_results_dir,
learning_rate=0.0001,
is_lr_scheduled=True,
max_epoch=1,
save_epochs=[10, 20, 30]):
# Measure execution time
train_start = time.time()
start_time = strftime('SSD__%dth_%H:%M_', gmtime())
# Define the Net
print('num_class: ', len(class_labels))
print('class_labels: ', class_labels)
ssd_net = SSD(len(class_labels))
# Set the parameter defined in the net to GPU
net = ssd_net
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.backends.cudnn.benchmark = True
net.cuda()
# Define the loss
center_var = 0.1
size_var = 0.2
criterion = MultiboxLoss([center_var, center_var, size_var, size_var],
iou_threshold=0.5,
neg_pos_ratio=3.0)
# Define Optimizer
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
# optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9,
# weight_decay=0.0005)
if is_lr_scheduled:
scheduler = MultiStepLR(optimizer,
milestones=[10, 30, 50, 70],
gamma=0.1)
# Train data
conf_losses = []
loc_losses = []
v_conf_losses = []
v_loc_losses = []
itr = 0
train_log = []
valid_log = []
for epoch_idx in range(0, max_epoch):
# decrease learning rate
if is_lr_scheduled:
scheduler.step()
print('\n\n===> lr: {}'.format(scheduler.get_lr()[0]))
# Save the trained network
if epoch_idx in save_epochs:
temp_file = start_time + 'epoch_{}'.format(epoch_idx)
net_state = net.state_dict() # serialize the instance
torch.save(net_state, lab_results_dir + temp_file + '__model.pth')
print('================> Temp file is created: ',
lab_results_dir + temp_file + '__model.pth')
# iterate the mini-batches:
for train_batch_idx, data in enumerate(train_loader):
train_images, train_labels, train_bboxes, prior_bbox = data
# Switch to train model
net.train()
# Forward
train_img = Variable(train_images.clone().cuda())
train_bbox = Variable(train_bboxes.clone().cuda())
train_label = Variable(train_labels.clone().cuda())
train_out_confs, train_out_locs = net.forward(train_img)
# locations(feature map base) -> bbox(center form)
train_out_bbox = loc2bbox(train_out_locs,
prior_bbox[0].unsqueeze(0))
# update the parameter gradients as zero
optimizer.zero_grad()
# Compute the loss
conf_loss, loc_loss = criterion.forward(train_out_confs,
train_out_bbox,
train_label, train_bbox)
train_loss = conf_loss + loc_loss
# Do the backward to compute the gradient flow
train_loss.backward()
# Update the parameters
optimizer.step()
conf_losses.append((itr, conf_loss))
loc_losses.append((itr, loc_loss))
itr += 1
if train_batch_idx % 20 == 0:
train_log_temp = '[Train]epoch: %d itr: %d Conf Loss: %.4f Loc Loss: %.4f' % (
epoch_idx, itr, conf_loss, loc_loss)
train_log += (train_log_temp + '\n')
print(train_log_temp)
if False: # check input tensor
image_s = train_images[0, :, :, :].cpu().numpy().astype(
np.float32).transpose().copy() # c , h, w -> h, w, c
image_s = ((image_s + 1) / 2)
bbox_cr_s = torch.cat([
train_bboxes[..., :2] - train_bboxes[..., 2:] / 2,
train_bboxes[..., :2] + train_bboxes[..., 2:] / 2
],
dim=-1)
bbox_prior_s = bbox_cr_s[0, :].cpu().numpy().astype(
np.float32).reshape(
(-1, 4)).copy() # First sample in batch
bbox_prior_s = (bbox_prior_s * 300)
label_prior_s = train_labels[0, :].cpu().numpy().astype(
np.float32).copy()
bbox_s = bbox_prior_s[label_prior_s > 0]
label_s = (label_prior_s[label_prior_s > 0]).astype(
np.uint8)
for idx in range(0, len(label_s)):
cv2.rectangle(image_s,
(bbox_s[idx][0], bbox_s[idx][1]),
(bbox_s[idx][2], bbox_s[idx][3]),
(255, 0, 0), 2)
cv2.putText(image_s, class_labels[label_s[idx]],
(bbox_s[idx][0], bbox_s[idx][1]),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 0, 0),
1, cv2.LINE_AA)
plt.imshow(image_s)
plt.show()
# validaton
if train_batch_idx % 200 == 0:
net.eval() # Evaluation mode
v_conf_subsum = torch.zeros(
1) # collect the validation losses for avg.
v_loc_subsum = torch.zeros(1)
v_itr_max = 5
for valid_itr, data in enumerate(valid_loader):
valid_image, valid_label, valid_bbox, prior_bbox = data
valid_image = Variable(valid_image.cuda())
valid_bbox = Variable(valid_bbox.cuda())
valid_label = Variable(valid_label.cuda())
# Forward and compute loss
with torch.no_grad(
): # make all grad flags to false!! ( Memory decrease)
valid_out_confs, valid_out_locs = net.forward(
valid_image)
valid_out_bbox = loc2bbox(
valid_out_locs,
prior_bbox[0].unsqueeze(0)) # loc -> bbox(center form)
valid_conf_loss, valid_loc_loss = criterion.forward(
valid_out_confs, valid_out_bbox, valid_label,
valid_bbox)
v_conf_subsum += valid_conf_loss
v_loc_subsum += valid_loc_loss
valid_itr += 1
if valid_itr > v_itr_max:
break
# avg. valid loss
v_conf_losses.append((itr, v_conf_subsum / v_itr_max))
v_loc_losses.append((itr, v_loc_subsum / v_itr_max))
valid_log_temp = '==>[Valid]epoch: %d itr: %d Conf Loss: %.4f Loc Loss: %.4f' % (
epoch_idx, itr, v_conf_subsum / v_itr_max,
v_loc_subsum / v_itr_max)
valid_log += (valid_log_temp + '\n')
print(valid_log_temp)
# Measure the time
train_end = time.time()
m, s = divmod(train_end - train_start, 60)
h, m = divmod(m, 60)
# Save the result
results_file_name = start_time + 'itr_{}'.format(itr)
train_data = {
'conf_losses': np.asarray(conf_losses),
'loc_losses': np.asarray(loc_losses),
'v_conf_losses': np.asarray(v_conf_losses),
'v_loc_losses': np.asarray(v_loc_losses),
'learning_rate': learning_rate,
'total_itr': itr,
'max_epoch': max_epoch,
'train_time': '%d:%02d:%02d' % (h, m, s)
}
torch.save(train_data, lab_results_dir + results_file_name + '.loss')
# Save the trained network
net_state = net.state_dict() # serialize the instance
torch.save(net_state, lab_results_dir + results_file_name + '__model.pth')
# Save the train/valid log
torch.save({'log': train_log},
lab_results_dir + results_file_name + '__train.log')
torch.save({'log': valid_log},
lab_results_dir + results_file_name + '__valid.log')
return lab_results_dir + results_file_name
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)
real_bounding_box = torch.squeeze(real_bounding_box)
class_list, sel_box = nms_bbox(real_bounding_box,
conf,
overlap_threshold=0.5,
prob_threshold=0.4)
# print(len(sel_box))
img = imgs[0].permute(1, 2, 0).contiguous()
img = Image.fromarray(np.uint8(img * 128 + 127))
sel_box = np.array(sel_box)
loc_targets = torch.squeeze(loc2bbox(loc_targets, prior)).numpy()
fig, ax = plt.subplots(1)
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:
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]),
