def visual(clist, cdir, num_pts):
if not cdir.exists(): os.makedirs(str(cdir))
shape = 256
transform = transforms.Compose(
[transforms.PreCrop(0.2),
transforms.TrainScale2WH((shape, shape))])
data = datasets.GeneralDataset(transform, 2, 1, 'gaussian', 'test')
data.load_list(clist, num_pts, True)
for i, tempx in enumerate(data):
image = tempx[0]
heats = models.variable2np(tempx[1]).transpose(1, 2, 0)
xheat = generate_color_from_heatmaps(heats, index=-1)
xheat = PIL.Image.fromarray(np.uint8(xheat * 255))
cimage = overlap_two_pil_image(image, xheat)
basename = osp.basename(data.datas[i]).split('.')[0]
basename = str(cdir) + '/' + basename + '-{:}.jpg'
image.save(basename.format('ori'))
xheat.save(basename.format('heat'))
cimage.save(basename.format('over'))
if i % PRINT_GAP == 0:
print('--->>> process the {:4d}/{:4d}-th image'.format(
i, len(data)))
def evaluation_image(eval_loader, net, log, save_path, opt):
if not osp.isdir(save_path): os.makedirs(save_path)
if opt.debug_save:
debug_save_dir = osp.join(save_path, 'debug-eval')
if not os.path.isdir(debug_save_dir): os.makedirs(debug_save_dir)
else: debug_save_dir = None
print_log(
' ==>start evaluation image dataset, save into {} with the error bar : {}, using the last stage, DEBUG SAVE DIR : {}'
.format(save_path, opt.error_bar, debug_save_dir), log)
# switch to eval mode
net.eval()
# calculate the time
batch_time, end = AverageMeter(), time.time()
# save the evaluation results
eval_meta = Eval_Meta()
# save the number of points
num_pts, scale_eval = opt.num_pts, opt.scale_eval
for i, (inputs, target, mask, points, image_index, label_sign,
ori_size) in enumerate(eval_loader):
# inputs : Batch, Squence, Channel, Height, Width
target = target.cuda(async=True)
mask = mask.cuda(async=True)
# forward, batch_locs [1 x points] [batch, 2]
batch_cpms, batch_locs, batch_scos, generated = net(inputs)
assert batch_locs.size(0) == batch_scos.size(0) and batch_locs.size(
0) == inputs.size(0)
assert batch_locs.size(1) == num_pts + 1 and batch_scos.size(
1) == num_pts + 1
assert batch_locs.size(2) == 2 and len(batch_scos.size()) == 2
np_batch_locs, np_batch_scos = variable2np(batch_locs), variable2np(
batch_scos)
image_index = np.squeeze(variable2np(image_index)).tolist()
if isinstance(image_index, numbers.Number): image_index = [image_index]
sign_list = variable2np(label_sign).astype('bool').squeeze(1).tolist()
# recover the ground-truth label
real_sizess = variable2np(ori_size)
for ibatch, imgidx in enumerate(image_index):
locations, scores = np_batch_locs[ibatch, :-1, :], np_batch_scos[
ibatch, :-1]
if len(scores.shape) == 1: scores = np.expand_dims(scores, axis=1)
xpoints = eval_loader.dataset.labels[imgidx].get_points()
assert real_sizess[ibatch, 0] > 0 and real_sizess[
ibatch,
1] > 0, 'The ibatch={}, imgidx={} is not right.'.format(
ibatch, imgidx)
scale_h, scale_w = real_sizess[ibatch, 0] * 1. / inputs.size(
-2), real_sizess[ibatch, 1] * 1. / inputs.size(-1)
locations[:,
0], locations[:,
1] = locations[:, 0] * scale_w + real_sizess[
ibatch,
2], locations[:,
1] * scale_h + real_sizess[
ibatch, 3]
assert xpoints.shape[1] == num_pts and locations.shape[
0] == num_pts and scores.shape[
0] == num_pts, 'The number of points is {} vs {} vs {} vs {}'.format(
num_pts, xpoints.shape, locations.shape, scores.shape)
# recover the original resolution
locations = np.concatenate((locations, scores), axis=1)
image_path = eval_loader.dataset.datas[imgidx]
face_size = eval_loader.dataset.face_sizes[imgidx]
image_name = osp.basename(image_path)
locations = locations.transpose(1, 0)
eval_meta.append(locations, xpoints, image_path, face_size)
if opt.error_bar is not None:
errorpath = osp.join(save_path, image_name)
save_error_image(image_path,
xpoints,
locations,
opt.error_bar,
errorpath,
radius=5,
color=(30, 255, 30),
rev_color=(255, 30, 30),
fontScale=10,
text_color=(255, 255, 255))
if opt.debug_save:
print_log(
'DEBUG --- > [{:03d}/{:03d}] '.format(i, len(eval_loader)),
log)
main_debug_save(debug_save_dir, eval_loader, image_index, inputs,
batch_locs, target, points, sign_list, batch_cpms,
generated, log)
# measure elapsed time
batch_time.update(time.time() - end)
need_hour, need_mins, need_secs = convert_secs2time(
batch_time.avg * (len(eval_loader) - i - 1))
end = time.time()
if i % opt.print_freq_eval == 0 or i + 1 == len(eval_loader):
print_log(
' Evaluation: [{:03d}/{:03d}] Time {:6.3f} ({:6.3f}) Need [{:02d}:{:02d}:{:02d}]'
.format(i, len(eval_loader), batch_time.val, batch_time.avg,
need_hour, need_mins, need_secs) +
' locations.shape={}'.format(np_batch_locs.shape), log)
return eval_meta
def main_debug_save(debug_save_dir, loader, image_index, input_vars, batch_locs, target, points, sign_list, batch_cpms, generations, log):
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
batch_size, num_pts = batch_locs.size(0), batch_locs.size(1)
assert batch_size == len(image_index)
print_log('Save-dir : {} : check {}'.format(debug_save_dir, image_index), log)
for idx, imgid in enumerate(image_index):
basename = osp.basename( loader.dataset.datas[imgid] )
basename = basename.split('.')[0]
print_log('Start visualization done for [{:03d}/{:03d}] : {:5} | {:}'.format(idx, len(image_index), imgid, loader.dataset.datas[imgid]), log)
## Save all images
images = [input_vars[idx], generations[0][idx], generations[1][idx]]
_pil_images = []
for iG, xinput in enumerate(images):
xinput = xinput.clone()
Xinput = []
for t, m, s in zip(xinput, mean, std):
t = torch.mul(t, s)
t = torch.add(t, m)
Xinput.append( t )
xinput = torch.stack(Xinput)
if xinput.is_cuda: xinput = xinput.cpu()
image = transforms.ToPILImage()(xinput.data)
debug_save_path = os.path.join(debug_save_dir, '{}-G{}.png'.format(basename, iG))
image.save(debug_save_path)
_pil_images.append( image )
"""
debug_loc = []
for ipts in range(num_pts):
temploc = models.variable2np( batch_locs[idx][ipts] )
debug_loc.append( temploc )
debug_loc = np.array( debug_loc )
debug_save_path = os.path.join(debug_save_dir, '{}-ans-points.png'.format(basename))
pimage = draw_image_with_pts(image, debug_loc.transpose(1,0), radius=1, linewidth=1, fontScale=12, window=None)
pimage.save(debug_save_path)
"""
image = _pil_images[0]
debug_save_path = os.path.join(debug_save_dir, '{}-GG.png'.format(basename))
overlap_two_pil_image(_pil_images[1], _pil_images[2]).save(debug_save_path)
# save the back ground heatmap
for icpm in range(len(batch_cpms)):
cpms = batch_cpms[icpm][idx]
xtarget = models.variable2np( cpms )
xtarget = xtarget.transpose(1,2,0)
xheatmaps = generate_color_from_heatmaps(xtarget, index=-1)
cimage = PIL.Image.fromarray(np.uint8(xheatmaps*255))
cimage = overlap_two_pil_image(image, cimage)
debug_save_path = os.path.join(debug_save_dir, '{:}-BG-{}.png'.format(basename, icpm))
cimage.save(debug_save_path)
xheatmaps = generate_color_from_heatmaps(xtarget, index=0)
cimage = PIL.Image.fromarray(np.uint8(xheatmaps*255))
cimage = overlap_two_pil_image(image, cimage)
debug_save_path = os.path.join(debug_save_dir, '{:}-B0-{}.png'.format(basename, icpm))
cimage.save(debug_save_path)
# save the ground truth heatmap
if sign_list[idx] and False:
xtarget = models.variable2np( target[idx] )
xtarget = xtarget.transpose(1,2,0)
xheatmaps = generate_color_from_heatmaps(xtarget)
all_images = []
for pid in range(len(xheatmaps)):
cimage = PIL.Image.fromarray(np.uint8(xheatmaps[pid]*255))
cimage = overlap_two_pil_image(center_img, cimage)
all_images.append( cimage )
debug_save_path = os.path.join(debug_save_dir, '{}-gt-heatmap.png'.format(basename))
all_images = merge_images(all_images, 10)
all_images.save( debug_save_path )
# debug save for points
cpoints = models.variable2np( points[idx] )
debug_save_path = os.path.join(debug_save_dir, '{}-gt-points.png'.format(basename))
point_image = draw_image_with_pts(center_img, cpoints.transpose(1,0), radius=1, linewidth=1, fontScale=12, window=lk_window)
point_image.save(debug_save_path)
print_log('Calculate gt-visualization done for [{:03d}/{:03d}] : {:5} | {:}'.format(idx, len(image_index), imgid, loader.dataset.datas[imgid]), log)
else:
print_log('Skip the ground truth heatmap debug saving', log)
def train(train_loader, net, criterion, optimizer, epoch, opt, log):
batch_time = AverageMeter()
data_time = AverageMeter()
forward_time = AverageMeter()
visible_points = AverageMeter()
losses = AverageMeter()
# switch to train mode
net.module.set_mode('train')
criterion.train()
if opt.debug_save:
debug_save_dir = os.path.join(opt.save_path, 'debug-train-{:02d}'.format(epoch))
if not os.path.isdir(debug_save_dir): os.makedirs(debug_save_dir)
end = time.time()
for i, (inputs, target, mask, points, image_index, label_sign, _) in enumerate(train_loader):
# inputs : Batch, Squence, Channel, Height, Width
# data prepare
target = target.cuda(async=True)
# get the real mask
mask.masked_scatter_((1-label_sign).unsqueeze(-1).unsqueeze(-1), torch.ByteTensor(mask.size()).zero_())
mask_var = mask.cuda(async=True)
batch_size, num_pts = inputs.size(0), mask.size(1)-1
image_index = variable2np(image_index).squeeze(1).tolist()
# check the label indicator, whether is has annotation or not
sign_list = variable2np(label_sign).astype('bool').squeeze(1).tolist()
# measure data loading time
data_time.update(time.time() - end)
cvisible_points = torch.sum(mask[:,:-1,:,:]) * 1. / batch_size
visible_points.update(cvisible_points, batch_size)
# batch_cpms is a list for CPM stage-predictions, each element should be [Batch, Squence, C, H, W]
# batch_locs and batch_scos are two sequence-list of point-list, each element is [Batch, 2] / [Batch, 1]
# batch_next and batch_back are two sequence-list of point-list, each element is [Batch, 2] / [Batch, 2]
batch_cpms, batch_locs, batch_scos, generated = net(inputs)
forward_time.update(time.time() - end)
total_labeled_cpm = int(np.sum(sign_list))
# collect all cpm stages for the middle frame
cpm_loss, each_stage_loss_values = compute_stage_loss(criterion, target, batch_cpms, mask_var, total_labeled_cpm, opt.weight_of_idt)
# measure accuracy and record loss
losses.update(cpm_loss.item(), batch_size)
# compute gradient and do SGD step
optimizer.zero_grad()
cpm_loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
need_hour, need_mins, need_secs = convert_secs2time(batch_time.avg * (len(train_loader)-i-1))
last_time = '[{:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)
end = time.time()
if i % opt.print_freq == 0 or i+1 == len(train_loader):
print_log(' Epoch: [{:03d}/{:03d}][{:03d}/{:03d}] '
'Time {batch_time.val:5.2f} ({batch_time.avg:5.2f}) '
'Data {data_time.val:5.2f} ({data_time.avg:5.2f}) '
'Forward {forward_time.val:5.2f} ({forward_time.avg:5.2f}) '
'Loss {loss.val:6.3f} ({loss.avg:6.3f}) '.format(
epoch, opt.epochs, i, len(train_loader), batch_time=batch_time,
data_time=data_time, forward_time=forward_time, loss=losses)
+ last_time + show_stage_loss(each_stage_loss_values) \
+ ' In={} Tar={} Mask={}'.format(convert_size2str(inputs.size()), convert_size2str(target.size()), convert_size2str(mask_var.size())) \
+ ' Vis-PTS : {:2d} ({:.1f})'.format(int(visible_points.val), visible_points.avg), log)
# Just for debug and show the intermediate results.
if opt.debug_save:
print_log('DEBUG --- > [{:03d}/{:03d}] '.format(i, len(train_loader)), log)
main_debug_save(debug_save_dir, train_loader, image_index, inputs, batch_locs, target, points, sign_list, batch_cpms, generated, log)
return losses.avg