def main(argv):
(opts, args) = parser.parse_args(argv)
# Load experiment setting
assert isinstance(opts, object)
config = NetConfig(opts.config)
batch_size = config.hyperparameters['batch_size']
max_iterations = config.hyperparameters['max_iterations']
train_loader_a = get_data_loader(config.datasets['train_a'], batch_size)
train_loader_b = get_data_loader(config.datasets['train_b'], batch_size)
cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters['trainer']
local_dict = locals()
exec(cmd,globals(),local_dict)
trainer = local_dict['trainer']
# Check if resume training
iterations = 0
if opts.resume == 1:
iterations = trainer.resume(config.snapshot_prefix)
trainer.cuda(opts.gpu)
######################################################################################################################
# Setup logger and repare image outputs
train_writer = tensorboard.FileWriter("%s/%s" % (opts.log,os.path.splitext(os.path.basename(opts.config))[0]))
image_directory, snapshot_directory = prepare_snapshot_and_image_folder(config.snapshot_prefix, iterations, config.image_save_iterations)
for ep in range(0, MAX_EPOCHS):
for it, (images_a, images_b) in enumerate(izip(train_loader_a,train_loader_b)):
if images_a.size(0) != batch_size or images_b.size(0) != batch_size:
continue
images_a = Variable(images_a.cuda(opts.gpu))
images_b = Variable(images_b.cuda(opts.gpu))
# Main training code
trainer.dis_update(images_a, images_b, config.hyperparameters)
image_outputs = trainer.gen_update(images_a, images_b, config.hyperparameters)
assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs)
# Dump training stats in log file
if (iterations+1) % config.display == 0:
write_loss(iterations, max_iterations, trainer, train_writer)
if (iterations+1) % config.image_save_iterations == 0:
img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1)
write_html(snapshot_directory + "/index.html", iterations + 1, config.image_save_iterations, image_directory)
elif (iterations + 1) % config.image_display_iterations == 0:
img_filename = '%s/gen.jpg' % (image_directory)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1)
# Save network weights
if (iterations+1) % config.snapshot_save_iterations == 0:
trainer.save(config.snapshot_prefix, iterations)
iterations += 1
if iterations >= max_iterations:
return
def main(argv):
(opts, args) = parser.parse_args(argv)
if 'estimate' in opts.mode:
mode_idx = int(opts.mode[-1])
global colorPlatte, bones, Evaluation
if 'nyu' in opts.config:
colorPlatte = utils.util.nyuColorIdx
bones = utils.util.nyuBones
Evaluation = NYUHandposeEvaluation
elif 'icvl' in opts.config:
colorPlatte = utils.util.icvlColorIdx
bones = utils.util.icvlBones
Evaluation = ICVLHandposeEvaluation
# Load experiment setting
assert isinstance(opts, object)
config = NetConfig(opts.config)
batch_size = config.hyperparameters['batch_size'] if 'estimate' in opts.mode else 1
test_batch_size = batch_size * 32
max_iterations = config.hyperparameters['max_iterations']
frac = opts.frac
dataset_a = get_dataset(config.datasets['train_a'])
dataset_b = get_dataset(config.datasets['train_b'])
dataset_test = get_dataset(config.datasets['test_b'])
train_loader_a = get_data_loader(dataset_a, batch_size, shuffle=True)
train_loader_b = get_data_loader(dataset_b, batch_size, shuffle=True)
test_loader_real = get_data_loader(dataset_test, test_batch_size, shuffle=False)
cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters['trainer']
local_dict = locals()
exec(cmd,globals(),local_dict)
trainer = local_dict['trainer']
di_a = dataset_a.di
di_b = dataset_b.di
# Check if resume training
iterations = 0
if opts.resume == 1:
iterations = trainer.resume(config.snapshot_prefix, idx=-1, load_opt=True)
for i in range(iterations//1000):
trainer.dis_sch.step()
trainer.gen_sch.step()
trainer.cuda(opts.gpu)
print('using %.2f percent of the labeled real data' % frac)
try:
if 'estimate' in opts.mode and (mode_idx == 3 or mode_idx == 4):
trainer.load_vae(config.snapshot_prefix, 2+frac)
else:
trainer.load_vae(config.snapshot_prefix, frac)
except:
print('Failed to load the parameters of vae')
if 'estimate' in opts.mode:
if opts.idx != 0:
trainer.resume(config.snapshot_prefix, idx=opts.idx, est=mode_idx==5)
if frac > 0. and frac < 1.:
dataset_b.set_nmax(frac)
#trainer.dis.freeze_layers()
###############################################################################################
# Setup logger and repare image outputs
train_writer = tensorboardX.FileWriter("%s/%s" % (opts.log,os.path.splitext(os.path.basename(opts.config))[0]))
image_directory, snapshot_directory = prepare_snapshot_and_image_folder(config.snapshot_prefix, iterations, config.image_save_iterations)
best_err, best_acc = 100., 0.
start_time = time.time()
for ep in range(0, MAX_EPOCHS):
for it, ((images_a, labels_a, com_a, M_a, cube_a, _), (images_b,labels_b, com_b, M_b, cube_b, _)) in \
enumerate(izip(train_loader_a,train_loader_b)):
if images_a.size(0) != batch_size or images_b.size(0) != batch_size:
continue
images_a = Variable(images_a.cuda(opts.gpu))
images_b = Variable(images_b.cuda(opts.gpu))
labels_a = Variable(labels_a.cuda(opts.gpu))
labels_b = Variable(labels_b.cuda(opts.gpu))
com_a = Variable(com_a.cuda(opts.gpu))
com_b = Variable(com_b.cuda(opts.gpu))
trainer.dis.train()
if opts.mode == 'pretrain':
if (iterations+1) % 1000 == 0:
trainer.dis_sch.step()
trainer.gen_sch.step()
print('lr %.8f' % trainer.dis_sch.get_lr()[0])
trainer.dis_update(images_a, labels_a, images_b, labels_b, com_a, com_b, config.hyperparameters)
image_outputs = trainer.gen_update(images_a, labels_a, images_b, labels_b, config.hyperparameters)
assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs)
else:
if (iterations+1) % 100 == 0:
trainer.dis_sch.step()
image_outputs = trainer.post_update(images_a, labels_a, images_b, labels_b,com_a,com_b, mode_idx, config.hyperparameters)
assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs)
# Dump training stats in log file
if (iterations+1) % config.display == 0:
elapsed_time = time.time() - start_time
write_loss(iterations, max_iterations, trainer, train_writer, elapsed_time)
start_time = time.time()
if (iterations + 1) % config.image_display_iterations == 0:
img_filename = '%s/gen.jpg' % (image_directory)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1)
if (iterations+1) % config.image_save_iterations == 0:
if opts.mode == 'pretrain':# and (iterations+1) % (2*config.image_save_iterations) != 0:
img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1)
write_html(snapshot_directory + "/index.html", iterations + 1, \
config.image_save_iterations, image_directory)
else:
trainer.dis.eval()
score, maxerr = 0, 0
num_samples = 0
maxJntError = []
meanJntError = 0
img2sav = None
gt3D = []
joints = []
joints_imgcord = []
codec = cv2.VideoWriter_fourcc(*'XVID')
vid = cv2.VideoWriter(os.path.join(image_directory,'gen.avi'), codec, 25, (128*2,128))
for tit, (test_images_b, test_labels_b, com_b, trans_b, cube_b, fn) in enumerate(test_loader_real):
test_images_b = Variable(test_images_b.cuda(opts.gpu))
test_labels_b = Variable(test_labels_b.cuda(opts.gpu))
if mode_idx == 0:
pred_pose, pred_post, _ = trainer.dis.regress_a(test_images_b)
else:
pred_pose, pred_post, _ = trainer.dis.regress_b(test_images_b)
if True:
pred_pose = trainer.vae.decode(pred_post)
n = test_labels_b.size(0)
gt_pose = test_labels_b.data.cpu().numpy().reshape((n,-1, 3))
pr_pose = pred_pose.data.cpu().numpy().reshape((n,-1, 3))
if tit < 20:
for i in range(0, n, 4):
real_img = visPair(di_b, test_images_b[i].data.cpu().numpy(), gt_pose[i].reshape((-1)), \
trans_b[i].numpy(), com_b[i].numpy(), cube_b[i].numpy(), 50.0)
est_img = visPair(di_b, test_images_b[i].data.cpu().numpy(), pr_pose[i].reshape((-1)), \
trans_b[i].numpy(), com_b[i].numpy(), cube_b[i].numpy(), 50.0)
vid.write(np.hstack((real_img,est_img)).astype('uint8'))
both_img = np.vstack((real_img,est_img))
if True and tit < 8:
if img2sav is None:
img2sav = both_img
else:
img2sav = np.hstack((img2sav,both_img))
if 'nyu' in opts.config:
restrictedJointsEval = np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 27, 30, 31, 32])
gt_pose = gt_pose[:,restrictedJointsEval]
pr_pose = pr_pose[:,restrictedJointsEval]
for i in range(n):
gt3D.append(gt_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy())
joints.append(pr_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy())
joints_imgcord.append(di_b.joints3DToImg(pr_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy()))
score += meanJntError
num_samples += test_images_b.size(0)
cv2.imwrite(image_directory + '/_test.jpg', img2sav.astype('uint8'))
vid.release()
hpe = Evaluation(np.array(gt3D), np.array(joints))
mean_err = hpe.getMeanError()
over_40 = 100. * hpe.getNumFramesWithinMaxDist(40) / len(gt3D)
best_err = np.minimum(best_err, mean_err)
best_acc = np.maximum(best_acc, over_40)
print("------------ Mean err: {:.4f} ({:.4f}) mm, Max over 40mm: {:.2f} ({:.2f}) %".format(mean_err, best_err, over_40, best_acc))
# Save network weights
if (iterations+1) % config.snapshot_save_iterations == 0:
if opts.mode == 'pretrain':
trainer.save(config.snapshot_prefix, iterations)
elif 'estimate' in opts.mode:
trainer.save(config.snapshot_prefix+'_est', iterations)
iterations += 1
if iterations >= max_iterations:
return
def main(argv):
(opts, args) = parser.parse_args(argv)
# Load experiment setting
assert isinstance(opts, object)
config = NetConfig(opts.config)
batch_size = config.hyperparameters['batch_size']
max_iterations = config.hyperparameters['max_iterations']
train_loader_a = get_data_loader(config.datasets['train_a'], batch_size)
train_loader_b = get_data_loader(config.datasets['train_b'], batch_size)
# Parse ROI parameters
roi = [int(val_str) for val_str in opts.roi.split(',')]
roi_x = roi[0]
roi_y = roi[1]
roi_w = roi[2]
roi_h = roi[3]
cmd1 = "trainer=%s(config.hyperparameters)" % config.hyperparameters[
'trainer']
cmd2 = "roi_trainer=%s(config.hyperparameters)" % config.hyperparameters[
'trainer']
local_dict = locals()
exec(cmd1, globals(), local_dict)
trainer = local_dict['trainer']
exec(cmd2, globals(), local_dict)
roi_trainer = local_dict['roi_trainer']
# Check if resume training
iterations = 0
if opts.resume == 1:
iterations = trainer.resume(config.snapshot_prefix)
roi_trainer.resume(config.snapshot_prefix)
trainer.cuda(opts.gpu)
roi_trainer.cuda(opts.gpu)
######################################################################################################################
# Setup logger and repare image outputs
train_writer = tensorboard.FileWriter(
"%s/%s" %
(opts.log, os.path.splitext(os.path.basename(opts.config))[0]))
image_directory, snapshot_directory = prepare_snapshot_and_image_folder(
config.snapshot_prefix, iterations, config.image_save_iterations)
for ep in range(0, MAX_EPOCHS):
for it, (images_a,
images_b) in enumerate(izip(train_loader_a, train_loader_b)):
if images_a.size(0) != batch_size or images_b.size(
0) != batch_size:
continue
# Crop images according to ROI
roi_images_a = images_a[:, :, roi_y:roi_y + roi_h,
roi_x:roi_x + roi_w].clone()
roi_images_b = images_b[:, :, roi_y:roi_y + roi_h,
roi_x:roi_x + roi_w].clone()
roi_images_a = Variable(roi_images_a.cuda(opts.gpu))
roi_images_b = Variable(roi_images_b.cuda(opts.gpu))
images_a = Variable(images_a.cuda(opts.gpu))
images_b = Variable(images_b.cuda(opts.gpu))
# Main training code
trainer.dis_update(images_a, images_b, config.hyperparameters)
trainer.gen_update(images_a, images_b, config.hyperparameters)
# Training code for ROI
roi_trainer.dis_update(roi_images_a, roi_images_b,
config.hyperparameters)
roi_image_outputs = roi_trainer.gen_update(roi_images_a,
roi_images_b,
config.hyperparameters)
roi_assembled_images = roi_trainer.assemble_outputs(
roi_images_a, roi_images_b, roi_image_outputs)
# Paste ROI to original images to update generator
x_aa, x_ba, x_ab, x_bb, shared = trainer.gen(images_a, images_b)
x_ba_paste = x_ba.clone()
x_ab_paste = x_ab.clone()
x_ba_paste[:, :, roi_y:roi_y + roi_h,
roi_x:roi_x + roi_w] = roi_image_outputs[1].clone()
x_ab_paste[:, :, roi_y:roi_y + roi_h,
roi_x:roi_x + roi_w] = roi_image_outputs[2].clone()
trainer.gen.zero_grad()
image_outputs = trainer.gen_update_helper(images_a, images_b, x_aa,
x_ba_paste, x_ab_paste,
x_bb, shared,
config.hyperparameters)
assembled_images = trainer.assemble_outputs(
images_a, images_b, image_outputs)
# Dump training stats in log file
if (iterations + 1) % config.display == 0:
write_loss(iterations, max_iterations, trainer, train_writer)
if (iterations + 1) % config.image_save_iterations == 0:
img_filename = '%s/gen_%08d.jpg' % (image_directory,
iterations + 1)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5,
img_filename,
nrow=1)
img_filename = '%s/roi_gen_%08d.jpg' % (image_directory,
iterations + 1)
torchvision.utils.save_image(roi_assembled_images.data / 2 +
0.5,
img_filename,
nrow=1)
write_html(snapshot_directory + "/index.html", iterations + 1,
config.image_save_iterations, image_directory)
elif (iterations + 1) % config.image_display_iterations == 0:
img_filename = '%s/gen.jpg' % (image_directory)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5,
img_filename,
nrow=1)
img_filename = '%s/roi_gen.jpg' % (image_directory)
torchvision.utils.save_image(roi_assembled_images.data / 2 +
0.5,
img_filename,
nrow=1)
# Save network weights
if (iterations + 1) % config.snapshot_save_iterations == 0:
trainer.save(config.snapshot_prefix, iterations)
iterations += 1
if iterations >= max_iterations:
return
def main(argv):
(opts, args) = parser.parse_args(argv)
seed = 0
torch.cuda.manual_seed(seed) # Set seed for deterministic, is this enough?
torch.manual_seed(seed)
np.random.seed(seed=seed)
# Load experiment setting
assert isinstance(opts, object)
config = NetConfig(opts.config)
batch_size = config.hyperparameters['batch_size']
max_iterations = config.hyperparameters['max_iterations']
train_loader_a = get_data_loader(config.datasets['train_a'], batch_size)
train_loader_b = get_data_loader(config.datasets['train_b'], batch_size)
cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[
'trainer']
local_dict = locals()
exec(cmd, globals(), local_dict)
trainer = local_dict['trainer']
# Check if resume training
iterations = 0
if opts.resume == 1:
iterations = trainer.resume(config.snapshot_prefix)
trainer.cuda(opts.gpu)
######################################################################################################################
# Setup logger and repare image outputs
train_writer = tf.summary.FileWriter(
"%s/%s" %
(opts.log, os.path.splitext(os.path.basename(opts.config))[0]))
image_directory, snapshot_directory = prepare_snapshot_and_image_folder(
config.snapshot_prefix, iterations, config.image_save_iterations)
for ep in range(0, MAX_EPOCHS):
for it, (data_a,
data_b) in enumerate(izip(train_loader_a, train_loader_b)):
images_a = data_a['data']
labels_a = data_a.get("data_lab")
images_b = data_b['data']
labels_b = data_b.get("data_lab")
if images_a.size(0) != batch_size or images_b.size(
0) != batch_size:
continue
images_a = Variable(images_a.cuda(opts.gpu))
images_b = Variable(images_b.cuda(opts.gpu))
# Main training code
trainer.dis_update(images_a, images_b, config.hyperparameters)
image_outputs = trainer.gen_update(images_a, images_b,
config.hyperparameters,
labels_a, labels_b)
assembled_images = trainer.assemble_outputs(
images_a, images_b, image_outputs)
# Dump training stats in log file
if (iterations + 1) % config.display == 0:
write_loss(iterations, max_iterations, trainer, train_writer)
if (iterations + 1) % config.image_save_iterations == 0:
img_filename = '%s/gen_%08d.jpg' % (image_directory,
iterations + 1)
torchvision.utils.save_image(assembled_images.data,
img_filename,
nrow=1)
segm_image = None
if labels_a is not None:
_, enet_classes_ab = torch.max(image_outputs[7],
dim=1,
keepdim=False)
segm_image_ab = np.concatenate(
(np.squeeze(enet_classes_ab.data.cpu().numpy()),
np.squeeze(labels_a.cpu().numpy())),
axis=1)
segm_image = segm_image_ab
if labels_b is not None:
_, enet_classes_ba = torch.max(image_outputs[6],
dim=1,
keepdim=False)
segm_image_ba = np.concatenate(
(np.squeeze(enet_classes_ba.data.cpu().numpy()),
np.squeeze(labels_b.cpu().numpy())),
axis=1)
if labels_a is not None:
segm_image = np.concatenate(
(segm_image_ba, segm_image_ab), axis=0)
else:
segm_image = segm_image_ba
if segm_image is not None:
segm_filename = '%s/segm_cat_%08d.jpg' % (image_directory,
iterations + 1)
cv2.imwrite(segm_filename, segm_image)
write_html(snapshot_directory + "/index.html", iterations + 1,
config.image_save_iterations, image_directory)
elif (iterations + 1) % config.image_display_iterations == 0:
img_filename = '%s/gen.jpg' % (image_directory)
torchvision.utils.save_image(assembled_images.data,
img_filename,
nrow=1)
# Save network weights
if (iterations + 1) % config.snapshot_save_iterations == 0:
trainer.save(config.snapshot_prefix, iterations)
iterations += 1
if iterations >= max_iterations:
return
del images_a
del images_b
del image_outputs
del assembled_images
del labels_a
del labels_b
def main(argv):
(opts, args) = parser.parse_args(argv)
# Load experiment setting
assert isinstance(opts, object)
config = NetConfig(opts.config)
batch_size = config.hyperparameters['batch_size']
max_iterations = config.hyperparameters['max_iterations']
train_loader_a = get_data_loader(config.datasets['train_a'], batch_size)
train_loader_b = get_data_loader(config.datasets['train_b'], batch_size)
cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[
'trainer']
local_dict = locals()
exec(cmd, globals(), local_dict)
trainer = local_dict['trainer']
# Check if resume training
iterations = 0
if opts.resume == 1:
iterations = trainer.resume(config.snapshot_prefix)
trainer.cuda(opts.gpu)
######################################################################################################################
# Setup logger and repare image outputs
train_writer = tensorboard.FileWriter(
"%s/%s" %
(opts.log, os.path.splitext(os.path.basename(opts.config))[0]))
image_directory, snapshot_directory = prepare_snapshot_and_image_folder(
config.snapshot_prefix, iterations, config.image_save_iterations)
for ep in range(0, MAX_EPOCHS):
for it, (images_a,
images_b) in enumerate(izip(train_loader_a, train_loader_b)):
if images_a.size(0) != batch_size or images_b.size(
0) != batch_size:
continue
images_a = Variable(images_a.cuda(opts.gpu))
images_b = Variable(images_b.cuda(opts.gpu))
# Main training code
trainer.dis_update(images_a, images_b, config.hyperparameters)
image_outputs = trainer.gen_update(images_a, images_b,
config.hyperparameters)
assembled_images = trainer.assemble_outputs(
images_a, images_b, image_outputs)
# Dump training stats in log file
if (iterations + 1) % config.display == 0:
write_loss(iterations, max_iterations, trainer, train_writer)
if (iterations + 1) % config.image_save_iterations == 0:
img_filename = '%s/gen_%08d.jpg' % (image_directory,
iterations + 1)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5,
img_filename,
nrow=1)
write_html(snapshot_directory + "/index.html", iterations + 1,
config.image_save_iterations, image_directory)
elif (iterations + 1) % config.image_display_iterations == 0:
img_filename = '%s/gen.jpg' % (image_directory)
torchvision.utils.save_image(assembled_images.data / 2 + 0.5,
img_filename,
nrow=1)
# Save network weights
if (iterations + 1) % config.snapshot_save_iterations == 0:
trainer.save(config.snapshot_prefix, iterations)
iterations += 1
if iterations >= max_iterations:
return
def main(argv):
(opts, args) = parser.parse_args(argv)
global colorPlatte, bones, Evaluation
if 'nyu' in opts.config:
colorPlatte = utils.util.nyuColorIdx
bones = utils.util.nyuBones
Evaluation = NYUHandposeEvaluation
elif 'icvl' in opts.config:
colorPlatte = utils.util.icvlColorIdx
bones = utils.util.icvlBones
Evaluation = ICVLHandposeEvaluation
# Load experiment setting
assert isinstance(opts, object)
config = NetConfig(opts.config)
batch_size = config.hyperparameters['batch_size_pose']
max_iterations = 200000 #config.hyperparameters['max_iterations']
frac = opts.frac
dataset_a = get_dataset(config.datasets['train_a'])
dataset_b = get_dataset(config.datasets['train_b'])
dataset_test = get_dataset(config.datasets['test_b'])
train_loader_a = get_data_loader(dataset_a, batch_size, shuffle=True)
train_loader_b = get_data_loader(dataset_b, batch_size, shuffle=True)
test_loader_real = get_data_loader(dataset_test, 1, shuffle=True)
cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[
'trainer']
local_dict = locals()
exec(cmd, globals(), local_dict)
trainer = local_dict['trainer']
iterations = 0
trainer.cuda(opts.gpu)
dataset_a.pose_only = True
dataset_b.pose_only = True
if frac > 0. and frac < 1.:
dataset_b.set_nmax(frac)
di_a = dataset_a.di
di_b = dataset_b.di
dataset_a.sample_poses()
dataset_b.sample_poses()
###################################################################
# Setup logger and repare image outputs
train_writer = tensorboardX.FileWriter(
"%s/%s" %
(opts.log, os.path.splitext(os.path.basename(opts.config))[0]))
image_directory, snapshot_directory = prepare_snapshot_and_image_folder(
config.snapshot_prefix, iterations, config.image_save_iterations)
print('using %.2f percent of the labeled real data' % frac)
start_time = time.time()
for ep in range(0, MAX_EPOCHS):
for it, ((labels_a),
(labels_b)) in enumerate(izip(train_loader_a,
train_loader_b)):
if labels_a.size(0) != batch_size or labels_b.size(
0) != batch_size:
continue
labels_a = Variable(labels_a.cuda(opts.gpu))
labels_b = Variable(labels_b.cuda(opts.gpu))
labels = labels_a
if frac > 0.:
labels = torch.cat((labels_a, labels_b), 0)
if (iterations + 1) % 1000 == 0:
trainer.vae_sch.step()
recon_pose = trainer.vae_update(labels, config.hyperparameters)
# Dump training stats in log file
if (iterations + 1) % config.display == 0:
elapsed_time = time.time() - start_time
write_loss(iterations, max_iterations, trainer, train_writer,
elapsed_time)
start_time = time.time()
if (iterations + 1) % (10 * config.image_save_iterations) == 0:
if True:
score, maxerr = 0, 0
num_samples = 0
maxJntError = []
img2sav = None
gt3D = []
joints = []
for tit, (test_images_b, test_labels_b, com_b, trans_b,
cube_b, _) in enumerate(test_loader_real):
test_images_b = Variable(test_images_b.cuda(opts.gpu))
test_labels_b = Variable(test_labels_b.cuda(opts.gpu))
pred_pose = trainer.vae.decode(
trainer.vae.encode(test_labels_b)[1])
gt3D.append(test_labels_b.data.cpu().numpy().reshape((-1, 3))*(cube_b.numpy()[0]/2.) +\
com_b.numpy())
joints.append(pred_pose.data.cpu().numpy().reshape((-1, 3))*(cube_b.numpy()[0]/2.) +\
com_b.numpy())
if True and tit < 8:
real_img = visPair(di_b, test_images_b.data.cpu().numpy(), test_labels_b.data.cpu().numpy(), \
trans_b.numpy(), com_b.numpy(), cube_b.numpy(), 50.0)
est_img = visPair(di_b, test_images_b.data.cpu().numpy(), pred_pose.data.cpu().numpy(), \
trans_b.numpy(), com_b.numpy(), cube_b.numpy(), 50.0)
if img2sav is None:
img2sav = np.vstack((real_img, est_img))
else:
img2sav = np.hstack(
(img2sav, np.vstack((real_img, est_img))))
num_samples += test_images_b.size(0)
cv2.imwrite(image_directory + '/_test.jpg',
img2sav.astype('uint8'))
#maxerr = Evaluation.plotError(maxJntError, image_directory + '/maxJntError.txt')
hpe = Evaluation(np.array(gt3D), np.array(joints))
print("Mean error: {}mm, max error: {}mm".format(
hpe.getMeanError(), hpe.getMaxError()))
# Save network weights
if (iterations + 1) % (4 * config.snapshot_save_iterations) == 0:
trainer.save_vae(config.snapshot_prefix, iterations, 2 + frac)
iterations += 1
if iterations >= max_iterations:
return