def train(train_loader, net, optimizer, epoch, visualizer, idx, opt):
# batch_time = AverageMeter()
# data_time = AverageMeter()
losses = AverageMeter()
pckhs = AverageMeter()
pckhs_origin_res = AverageMeter()
# switch to train mode
net.train()
# end = time.time()
for i, (img, heatmap, c, s, r, grnd_pts,
normalizer) in enumerate(train_loader):
# """measure data loading time"""
# data_time.update(time.time() - end)
# input and groundtruth
# img_var = torch.autograd.Variable(img)
img = img.cuda(non_blocking=True)
heatmap = heatmap.cuda(non_blocking=True)
# target_var = torch.autograd.Variable(heatmap)
# output and loss
# output1, output2 = net(img_var)
# loss = (output1 - target_var) ** 2 + (output2 - target_var) ** 2
output = net(img)
# exit()
# print(type(output))
# print(len(output))
loss = 0
for per_out in output:
tmp_loss = (per_out - heatmap)**2
loss = loss + tmp_loss.sum() / tmp_loss.numel()
# gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# """measure optimization time"""
# batch_time.update(time.time() - end)
# end = time.time()
# print log
losses.update(loss.item())
pckh = Evaluation.accuracy(output[-1].cpu(), heatmap.cpu(), idx)
pckhs.update(pckh[0])
pckh_origin_res = Evaluation.accuracy_origin_res(
output[-1].cpu(), c, s, [64, 64], grnd_pts, normalizer, r)
pckhs_origin_res.update(pckh_origin_res[0])
loss_dict = OrderedDict([('loss', losses.avg), ('pckh', pckhs.avg),
('pckh_origin_res', pckhs_origin_res.avg)])
if i % opt.print_freq == 0 or i == len(train_loader) - 1:
visualizer.print_log(epoch, i, len(train_loader), value1=loss_dict)
# if i == 1:
# break
return losses.avg, pckhs_origin_res.avg
def compute_separated_s_r_pckh(hg, img_list, c_list, s_list, r_list,
grnd_pts_list, grnd_heatmap_list,
normalizer_list):
assert len(img_list) == 2
pckh_list = []
for k in range(0, 2):
img_var = torch.autograd.Variable(img_list[k])
out_reg = hg(img_var)
tmp_pckh = Evaluation.per_person_pckh(out_reg[-1].data.cpu(),
grnd_heatmap_list[k], c_list[k],
s_list[k], [64, 64],
grnd_pts_list[k],
normalizer_list[k], r_list[k])
pckh_list.append(tmp_pckh)
return pckh_list
def validate(val_loader, net, epoch, visualizer, idx, num_classes):
batch_time = AverageMeter()
losses_det = AverageMeter()
losses = AverageMeter()
pckhs = AverageMeter()
pckhs_origin_res = AverageMeter()
img_batch_list = []
pts_batch_list = []
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), num_classes, 2)
# switch to evaluate mode
net.eval()
end = time.time()
for i, (img, heatmap, center, scale, rot, grnd_pts, normalizer,
index) in enumerate(val_loader):
# input and groundtruth
input_var = torch.autograd.Variable(img, volatile=True)
heatmap = heatmap.cuda(async=True)
target_var = torch.autograd.Variable(heatmap)
# output and loss
#output1, output2 = net(input_var)
#loss = (output1 - target_var) ** 2 + (output2 - target_var) ** 2
output1 = net(input_var)
loss = 0
for per_out in output1:
tmp_loss = (per_out - target_var)**2
loss = loss + tmp_loss.sum() / tmp_loss.numel()
# flipping the image
img_flip = img.numpy()[:, :, :, ::-1].copy()
img_flip = torch.from_numpy(img_flip)
input_var = torch.autograd.Variable(img_flip, volatile=True)
#output11, output22 = net(input_var)
output2 = net(input_var)
output2 = HumanAug.flip_channels(output2[-1].cpu().data)
output2 = HumanAug.shuffle_channels_for_horizontal_flipping(output2)
output = (output1[-1].cpu().data + output2) / 2
# calculate measure
# pred_pts = HumanPts.heatmap2pts(output) # b x L x 2
# pts = HumanPts.heatmap2pts(target_var.cpu().data)
# pckh = HumanAcc.approx_PCKh(pred_pts, pts, idx, heatmap.size(3)) # b -> 1
pckh = Evaluation.accuracy(output, target_var.data.cpu(), idx)
pckhs.update(pckh[0])
pckh_origin_res = Evaluation.accuracy_origin_res(
output, center, scale, [64, 64], grnd_pts, normalizer, rot)
pckhs_origin_res.update(pckh_origin_res[0])
"""measure elapsed time"""
batch_time.update(time.time() - end)
end = time.time()
# print(log)
losses.update(loss.data[0])
loss_dict = OrderedDict([('loss', losses.avg), ('pckh', pckhs.avg),
('pckh_origin_res', pckhs_origin_res.avg)])
visualizer.print_log(epoch, i, len(val_loader), value1=loss_dict)
# img_batch_list.append(img)
# pts_batch_list.append(pred_pts*4.)
# preds = Evaluation.final_preds(output, meta['center'], meta['scale'], [64, 64])
# for n in range(output.size(0)):
# predictions[meta['index'][n], :, :] = preds[n, :, :]
preds = Evaluation.final_preds(output, center, scale, [64, 64], rot)
for n in range(output.size(0)):
predictions[index[n], :, :] = preds[n, :, :]
# if i == 2:
# break
# return losses.avg, pckhs.avg, img_batch_list, pts_batch_list
return losses.avg, pckhs_origin_res.avg, predictions
def train(train_loader, net, optimizer, epoch, visualizer, idx, opt):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
pckhs = AverageMeter()
pckhs_origin_res = AverageMeter()
# switch to train mode
net.train()
end = time.time()
for i, (img, heatmap, c, s, r, grnd_pts,
normalizer) in enumerate(train_loader):
# print('r: ', r)
# print('s: ', s)
# print('grnd_pts: ', pts)
# print('pts_aug_back: ', pts_aug_back)
# exit()
"""measure data loading time"""
data_time.update(time.time() - end)
# input and groundtruth
img_var = torch.autograd.Variable(img)
heatmap = heatmap.cuda(async=True)
target_var = torch.autograd.Variable(heatmap)
# output and loss
#output1, output2 = net(img_var)
#loss = (output1 - target_var) ** 2 + (output2 - target_var) ** 2
output = net(img_var)
# print(type(output))
# print(len(output))
loss = 0
for per_out in output:
tmp_loss = (per_out - target_var)**2
loss = loss + tmp_loss.sum() / tmp_loss.numel()
# loss = criterion(per_out, target_var)
# gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
"""measure optimization time"""
batch_time.update(time.time() - end)
end = time.time()
# print(log)
losses.update(loss.data[0])
# pred_pts = HumanPts.heatmap2pts(output[-1].cpu().data) # b x L x 2
# pts = HumanPts.heatmap2pts(target_var.cpu().data)
# pckh = HumanAcc.approx_PCKh(pred_pts, pts, idx, heatmap.size(3))
pckh = Evaluation.accuracy(output[-1].data.cpu(),
target_var.data.cpu(), idx)
pckhs.update(pckh[0])
pckh_origin_res = Evaluation.accuracy_origin_res(
output[-1].data.cpu(), c, s, [64, 64], grnd_pts, normalizer, r)
pckhs_origin_res.update(pckh_origin_res[0])
loss_dict = OrderedDict([('loss', losses.avg), ('pckh', pckhs.avg),
('pckh_origin_res', pckhs_origin_res.avg)])
if i % opt.print_freq == 0 or i == len(train_loader) - 1:
visualizer.print_log(epoch, i, len(train_loader), value1=loss_dict)
# if i == 2:
# break
return losses.avg, pckhs_origin_res.avg
def collect_data(data_loader, dataset, hg, save_path):
rot_num = len(dataset.rotation_means)
# rotation_num = len(dataset.rotation_means)
print 'rot_num: ', rot_num
# print 'rotation_num: ', rotation_num
# scale_1_idx = scale_num / 2
# print 'scale_1_idx', scale_1_idx
grnd_distri_list = []
# index_list = []
counter = 0
hg.eval()
# multipliers = np.zeros(7)
# scale_factors = np.arange(0.7, 1.4, 0.1)
# for i, s in enumerate(scale_factors):
# multipliers[i] = s
# print 'multipliers: ', multipliers
# exit()
for i, (img_list, heatmap_list, center_list, scale_list, rot_list,
grnd_pts_list, normalizer_list, rot_idx, img_index, pts_aug_list)\
in enumerate(data_loader):
print '%d/%d' % (i, len(data_loader))
# pts_aug_back = Evaluation.transform_preds(pts_aug_list[0][0]+1, center_list[0][0],
# scale_list[0][0], [64, 64], rot_list[0][0])
# print 'grnd_pts: ', grnd_pts_list[0][0]
# print 'pts_aug_back: ', pts_aug_back
# exit()
# print 'rot_idx size ', rot_idx.size()
# exit()
# print img_index
# for j in range(0, len(img_list)):
# # print img_list[j].size()
# img_show = imutils.im_to_numpy(img_list[j][0])*255
# img_show = Image.fromarray(img_show.astype('uint8'), 'RGB')
# img_show.save('debug-images/%d.jpg' % j)
# exit()
# index_list.append(img_index)
# print img_index
# print(len(img_list))
# print(img_list[0].size())
# print img_list[0][0, 0, 0, 0], img_list[1][0, 0, 0, 0]
# print img_list[0][1, 0, 0, 0], img_list[1][1, 0, 0, 0]
# exit()
val_img_index = torch.arange(counter, counter + len(img_index)).long()
assert ((val_img_index - img_index).sum() == 0)
counter += len(img_index)
# print scale_ind.size()
# print scale_ind
# exit()
# print rotation_ind.size()
# print rotation_ind
# exit()
unique_num = rot_idx.size(0)
assert rot_idx.size(1) == rot_num
# print 'unique_num: ', unique_num
img = torch.cat(img_list, dim=0)
# print 'img size: ', img.size()
# print img[0, 0, 0, 0], img[2, 0, 0, 0]
# print img[1, 0, 0, 0], img[3, 0, 0, 0]
# exit()
pts_aug = torch.cat(pts_aug_list, dim=0)
heatmap = torch.cat(heatmap_list, dim=0)
# print 'heatmap size: ', heatmap.size()
center = torch.cat(center_list, dim=0)
# print 'center size: ', center.size()
# print 'center: ', center
# print 'center list 0: ', center_list[0]
# exit()
scale = torch.cat(scale_list, dim=0)
# print 'scale size: ', scale.size()
rotation = torch.cat(rot_list, dim=0)
grnd_pts = torch.cat(grnd_pts_list, dim=0)
# print 'grnd_pts size: ', grnd_pts.size()
normalizer = torch.cat(normalizer_list, dim=0)
# print 'normalizer size: ', normalizer.size()
# exit()
# batch_size = img.size(0)
# print 'batch_size: ', batch_size
# output and loss
img_var = torch.autograd.Variable(img, volatile=True)
out_reg = hg(img_var)
# loss = 0
# for per_out in out_reg:
# # print 'hg', counter
# # counter += 1
# # print per_out.size()
# per_out = per_out.data.cpu()
# loss = loss + (per_out - heatmap) ** 2
# # loss = loss + tmp_loss.sum() / tmp_loss.numel()
# # exit()
# # print 'loss type: ', type(loss)
# # print 'loss size: ', loss.size()
#
# elm_num = loss.numel() / batch_size
# # print 'elm_num: ', elm_num
# loss = loss.view(loss.size(0), -1).sum(1).div_(elm_num)
# loss = loss.squeeze().numpy()
# print 'loss: ', loss
pckhs = Evaluation.per_person_pckh(out_reg[-1].data.cpu(), heatmap,
center, scale, [64, 64], grnd_pts,
normalizer, rotation)
lost_pckhs = 1 - pckhs
# pckhs = pckhs.numpy()
# print 'pckhs: ', pckhs
# print 'pred_counts shape: ', pred_counts.shape
for j in range(0, unique_num):
tmp_pckhs = lost_pckhs[j::unique_num]
# tmp_loss = loss[j::unique_num]
# print 'tmp_loss:', tmp_loss
# print 'weighted tmp_loss:', tmp_loss * multipliers
# print 'tmp_pckh:', tmp_pckhs
# exit()
assert (tmp_pckhs.size(0) == rot_num)
if tmp_pckhs.sum() == 0:
print 'tmp_pckh: ', tmp_pckhs
print 'sum of tmp_pckh are zero. Setting equal probabilities ...'
tmp_distri = torch.ones(tmp_pckhs.size(0)) / tmp_pckhs.size(0)
elif (tmp_pckhs < 0).any():
print 'tmp_pckh: ', tmp_pckhs
print 'some of tmp_pckh is negative. error...'
exit()
else:
tmp_distri = tmp_pckhs.clone()
tmp_distri = tmp_distri / tmp_distri.sum()
# print 'tmp_distri: ', tmp_distri
grnd_distri_list.append(tmp_distri)
with open(save_path, 'a+') as log_file:
tmp_distri = tmp_distri.numpy()
np.savetxt(log_file,
tmp_distri.reshape(1, tmp_distri.shape[0]),
fmt='%.2f')
# assert grnd_scale < scale_num
# grnd_scale = torch.LongTensor([grnd_scale])
# grnd_scale_list.append(grnd_scale)
# if i == 0:
# break
return grnd_distri_list
def train_hg(train_loader, hg, optimizer_hg, agent_sr, epoch, visualizer, opt):
batch_time = AverageMeter()
data_time = AverageMeter()
losses_hg_regular = AverageMeter()
losses_hg_sr = AverageMeter()
losses_hg = AverageMeter()
pckhs_regular = AverageMeter()
pckhs_sr = AverageMeter()
pckhs = AverageMeter()
# switch mode
hg.train()
agent_sr.eval()
# flags = ['neck', 'skip1', 'skip2', 'skip3', 'skip4']
end = time.time()
counter = 0
# idx_pckh = [e for e in idx if e not in (6, 7, 8, 9, 12, 13)]
drop_count = {}
is_asn = False
for i, (img_std, img, heatmap, c, s, r, grnd_pts, normalizer,
img_index) in enumerate(train_loader):
"""measure data loading time"""
# print img_index
data_time.update(time.time() - end)
# batch_size = img.size(0)
# save_imgs(img_std, 'std-imgs')
# save_imgs(img, 'regular-aug-imgs')
# input and groundtruth
if i % 2 == 0:
print 'regular augmentation'
img_var = torch.autograd.Variable(img)
# pts = HumanPts.heatmap2pts(heatmap)
target_var = torch.autograd.Variable(heatmap.cuda(async=True),
requires_grad=False)
out_reg = hg(img_var)
loss_hg_regular = 0
for per_out in out_reg:
tmp_loss = (per_out - target_var)**2
loss_hg_regular = loss_hg_regular + tmp_loss.sum(
) / tmp_loss.numel()
optimizer_hg.zero_grad()
loss_hg_regular.backward()
optimizer_hg.step()
losses_hg_regular.update(loss_hg_regular.data[0])
losses_hg.update(loss_hg_regular.data[0])
pckh = Evaluation.accuracy_origin_res(out_reg[-1].data.cpu(), c, s,
[64, 64], grnd_pts,
normalizer, r)
pckhs_regular.update(pckh[0])
pckhs.update(pckh[0])
else:
print 'agent augmentation'
img_var = torch.autograd.Variable(img_std)
pred_scale_distri, pred_rotation_distri = hg(img_var,
agent_sr,
is_half_hg=True,
is_aug=True)
pred_scale_distri = softmax(pred_scale_distri)
pred_rotation_distri = softmax(pred_rotation_distri)
pred_scale_distri_numpy = pred_scale_distri.data.cpu().numpy()
pred_rotation_distri_numpy = pred_rotation_distri.data.cpu().numpy(
)
# print pred_scale_distri_numpy
# print pred_rotation_distri_numpy
# exit()
scale_index_list = []
rotation_index_list = []
for j in range(0, pred_scale_distri_numpy.shape[0]):
# print len(dataset.scale_means), pred_scale_distri_numpy[j]
tmp_scale_index = np.random.choice(
len(dataset.scale_means), 1,
p=pred_scale_distri_numpy[j])[0]
# print pred_scale_distri_numpy[j], np.sum(pred_scale_distri_numpy[j]), tmp_scale_index
tmp_rotation_index = np.random.choice(
len(dataset.rotation_means),
1,
p=pred_rotation_distri_numpy[j])[0]
# print pred_rotation_distri_numpy[j], np.sum(pred_rotation_distri_numpy[j]), tmp_rotation_index
scale_index_list.append(tmp_scale_index)
rotation_index_list.append(tmp_rotation_index)
# if j == 1:
# exit()
# exit()
img, heatmap, c, s, r,\
grnd_pts, normalizer = load_batch_data(scale_index_list,
rotation_index_list,
img_index, opt,
separate_s_r=False)
# save_imgs(img, 'agent-aug-imgs')
# exit()
img_var = torch.autograd.Variable(img)
target_var = torch.autograd.Variable(heatmap.cuda(async=True),
requires_grad=False)
out_reg = hg(img_var)
loss_hg_sr = 0
for per_out in out_reg:
tmp_loss = (per_out - target_var)**2
loss_hg_sr = loss_hg_sr + tmp_loss.sum() / tmp_loss.numel()
optimizer_hg.zero_grad()
loss_hg_sr.backward()
optimizer_hg.step()
losses_hg_sr.update(loss_hg_sr.data[0])
losses_hg.update(loss_hg_sr.data[0])
pckh = Evaluation.accuracy_origin_res(out_reg[-1].data.cpu(), c, s,
[64, 64], grnd_pts,
normalizer, r)
pckhs_sr.update(pckh[0])
pckhs.update(pckh[0])
loss_dict = OrderedDict([('loss_hg_regular', losses_hg_regular.avg),
('loss_hg_sr', losses_hg_sr.avg),
('loss_hg', losses_hg.avg),
('pckhs_regular', pckhs_regular.avg),
('pckhs_sr', pckhs_sr.avg),
('pckh', pckhs.avg)])
# else:
# loss_dict = OrderedDict([('loss_hg_normal', losses_normal_hg.avg),
# ('pckh', pckhs.avg)])
if i % opt.print_freq == 0 or i == len(train_loader) - 1:
visualizer.print_log(epoch, i, len(train_loader), value1=loss_dict)
# if i == 1:
# break
return losses_hg.avg, pckhs.avg
def validate(val_loader, net, epoch, visualizer, num_classes, flip_index):
batch_time = AverageMeter()
losses_det = AverageMeter()
losses = AverageMeter()
rmses0 = AverageMeter()
rmses1 = AverageMeter()
rmses2 = AverageMeter()
img_batch_list = []
pts_batch_list = []
predictions = torch.Tensor(val_loader.dataset.__len__(), num_classes, 2)
# switch to evaluate mode
net.eval()
end = time.time()
for i, (img, heatmap, pts, index, center, scale) in enumerate(val_loader):
# input and groundtruth
input_var = torch.autograd.Variable(img, volatile=True)
heatmap = heatmap.cuda(async=True)
target_var = torch.autograd.Variable(heatmap)
# output and loss
#output1, output2 = net(input_var)
#loss = (output1 - target_var) ** 2 + (output2 - target_var) ** 2
output1 = net(input_var)
loss = 0
for per_out in output1:
tmp_loss = (per_out - target_var)**2
loss = loss + tmp_loss.sum() / tmp_loss.numel()
# # flipping the image
# img_flip = img.numpy()[:, :, :, ::-1].copy()
# img_flip = torch.from_numpy(img_flip)
# input_var = torch.autograd.Variable(img_flip, volatile=True)
# #output11, output22 = net(input_var)
# output2 = net(input_var)
# output2 = HumanAug.flip_channels(output2[-1].cpu().data)
# output2 = HumanAug.shuffle_channels_for_horizontal_flipping(output2, flip_index)
# output = (output1[-1].cpu().data + output2) / 2
# calculate measure
output = output1[-1].data.cpu()
# pred_pts_0, pred_pts_1, pred_pts_2 = FaceAcc.heatmap2pts(output)
# pred_pts_origin_0 = pred_pts_0.clone()
# pred_pts_origin_1 = pred_pts_1.clone()
# pred_pts_origin_2 = pred_pts_2.clone()
# for j in range(pred_pts_0.size(0)):
# # print type(coords[i]), type(center[i]), type(scale[i])
# tmp_pts = HumanAug.TransformPts(pred_pts_0[j].numpy()-1,
# center[j].numpy(), scale[j].numpy(), 0, 64, 200, invert=1)
# pred_pts_origin_0[j] = torch.from_numpy(tmp_pts)
# pred_pts_origin_1[j] = Evaluation.transform_preds(pred_pts_1[j], center[j], scale[j], [64, 64])
# pred_pts_origin_2[j] = Evaluation.transform_preds(pred_pts_2[j], center[j], scale[j], [64, 64])
# rmse0 = np.sum(FaceAcc.per_image_rmse(pred_pts_origin_0.numpy(), pts.numpy())) / img.size(0)
# rmse1 = np.sum(FaceAcc.per_image_rmse(pred_pts_origin_1.numpy(), pts.numpy())) / img.size(0)
# rmse2 = np.sum(FaceAcc.per_image_rmse(pred_pts_origin_2.numpy(), pts.numpy())) / img.size(0)
# rmses0.update(rmse0, img.size(0))
# rmses1.update(rmse1, img.size(0))
# rmses2.update(rmse2, img.size(0))
preds = Evaluation.final_preds(output, center, scale, [64, 64])
rmse = np.sum(FaceAcc.per_image_rmse(preds.numpy(),
pts.numpy())) / img.size(0)
rmses2.update(rmse, img.size(0))
"""measure elapsed time"""
batch_time.update(time.time() - end)
end = time.time()
# print log
losses.update(loss.data[0])
loss_dict = OrderedDict([('loss', losses.avg), ('rmse', rmses2.avg)])
visualizer.print_log(epoch,
i,
len(val_loader),
batch_time.avg,
value1=loss_dict)
# preds = Evaluation.final_preds(output, center, scale, [64, 64])
for n in range(output.size(0)):
predictions[index[n], :, :] = preds[n, :, :]
return losses.avg, rmses2.avg, predictions