def main(args):
# initialize workspace
split_file, experiment_directory, experiment_directory_color = init_info(args.class_name)
instance_num = args.instance_id
# manually set a initial camera intrinsic parameters
img_h, img_w = 137, 137
intrinsic = np.array([[150., 0., 68.5], [0., 150., 68.5], [0., 0., 1.]])
# load codes
shape_code_fname = os.path.join(experiment_directory, 'LatentCodes', '{}.pth'.format(args.checkpoint))
shape_code = torch.load(shape_code_fname)['latent_codes'][instance_num].detach().cpu()
color_code_fname = os.path.join(experiment_directory, 'Color', 'LatentCodes', '{}.pth'.format(args.checkpoint_color))
color_code = torch.load(color_code_fname)['latent_codes'][instance_num].detach().cpu()
shape_code = shape_code.float().cuda()
color_code = color_code.float().cuda()
# load decoders
decoder = load_decoder(experiment_directory, args.checkpoint)
decoder = decoder.module.cuda()
decoder_color = load_decoder(experiment_directory, args.checkpoint_color, color_size=args.color_size, experiment_directory_color=experiment_directory_color)
decoder_color = decoder_color.module.cuda()
# generate camera views
img_hw = (img_h, img_w)
print('Image size: {0}.'. format(img_hw))
# [Note] To render high-quality consistent visualization, we set ray marching ratio back to 1.0 here. If speed up is needed, aggressive strategy can still be used.
output_shape = args.resolution
sdf_renderer = SDFRenderer(decoder, decoder_color, resize_intrinsic(intrinsic, output_shape / img_h), march_step=100, buffer_size=1, threshold=args.threshold, ray_marching_ratio=1.0)
camera_list = generate_camera_list(sdf_renderer)
# visualization settings
dist_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
prefix = os.path.join(dist_path, args.vis_folder, '{0}_{1}'.format(args.class_name, instance_num), 'vis')
basedir = os.path.dirname(prefix)
if not os.path.exists(basedir):
os.makedirs(basedir)
else:
import shutil
shutil.rmtree(basedir)
os.makedirs(basedir)
for idx, camera in enumerate(tqdm(camera_list)):
prefix_idx = prefix + '_{0:04d}'.format(idx)
demo_color_save_render_output(prefix_idx, sdf_renderer, shape_code, color_code, camera)
def train(args):
# initialize output_dir
if not os.path.exists(args.vis_folder):
os.makedirs(args.vis_folder)
#########################################################################
# load data
#########################################################################
mesh_data_dir, experiment_directory, split_file, synthetic_data_dir = init_info(
)
upper_loader = LoaderSingle(synthetic_data_dir, mesh_data_dir,
experiment_directory, split_file)
shape_md5, image_data, mesh_data, camera, depth = upper_loader[6]
img, _, normal, _ = image_data
gt_samples, norm_params = mesh_data
points_gt = np.array(gt_samples.vertices)
points_gt = (points_gt + norm_params['offset']) * norm_params['scale']
gt_pack = {}
gt_pack['depth'] = torch.from_numpy(depth).cuda()
gt_pack['normal'] = torch.from_numpy(normal).cuda()
gt_pack['silhouette'] = torch.from_numpy(depth < 1e5).type(
torch.uint8).cuda()
# # visualize gt
# cv2.imwrite('img.png', img)
# visualize_depth('test0.png', depth)
# with open('camera.pkl', 'wb') as f:
# pickle.dump(camera, f)
#########################################################################
# initialize tensor
#########################################################################
decoder = load_decoder(experiment_directory, args.checkpoint)
decoder = decoder.module.cuda()
evaluator = Evaluator(decoder)
latent_size = 256
std_ = 0.1
rand_tensor = torch.ones(1, latent_size).normal_(mean=0, std=std_)
if args.use_random_init:
latent_tensor = rand_tensor
else:
latent_code_dir = os.path.join(synthetic_data_dir, 'latent_codes',
'{0}.pth'.format(shape_md5))
latent_code = torch.load(latent_code_dir)
latent_tensor = latent_code[0].detach().cpu()
latent_size = latent_tensor.shape[-1]
if not args.use_gt_code:
latent_tensor = latent_tensor + rand_tensor
latent_tensor = latent_tensor.float().cuda()
if (not args.profile) and (not args.no_pretest):
points_tmp = evaluator.latent_vec_to_points(latent_tensor,
fname=os.path.join(
args.vis_folder,
'initial.ply'),
silent=True)
if points_tmp is None:
print(
'The current latent code does not correspond to a valid shape.'
)
dist = None
else:
dist = evaluator.compute_chamfer_distance(points_gt, points_tmp)
print('STD: {0:.3f}'.format(std_))
print('CHAMFER DISTANCE: {0:.3f}'.format(dist * 1000))
latent_tensor.requires_grad = True
optimizer_latent = torch.optim.Adam([latent_tensor], lr=args.lr)
#########################################################################
# optimization
#########################################################################
weight_dict = {}
weight_dict['w_depth'] = 10.0
weight_dict['w_normal'] = 5.0
weight_dict['w_mask_gt'] = 1.0
weight_dict['w_mask_out'] = 1.0
weight_dict['w_l2reg'] = 1.0
img_h, img_w = img.shape[0], img.shape[1]
img_hw = (img_h, img_w)
print('Image size: {0}.'.format(img_hw))
if args.visualize:
visualizer = Visualizer(img_hw)
visualizer.add_chamfer(dist)
else:
visualizer = None
# initialize renderer
if args.use_multiscale:
sdf_renderer = SDFRenderer(decoder,
camera.intrinsic,
img_hw=img_hw,
march_step=100,
buffer_size=1,
threshold=args.threshold,
ray_marching_ratio=args.ratio,
use_depth2normal=args.use_depth2normal)
sdf_renderer_1_2 = SDFRenderer(decoder,
downsize_camera_intrinsic(
camera.intrinsic, 2),
march_step=100,
buffer_size=3,
threshold=args.threshold,
use_depth2normal=args.use_depth2normal)
sdf_renderer_1_4 = SDFRenderer(decoder,
downsize_camera_intrinsic(
camera.intrinsic, 4),
march_step=100,
buffer_size=5,
threshold=args.threshold,
use_depth2normal=args.use_depth2normal)
renderer_list = [sdf_renderer, sdf_renderer_1_2, sdf_renderer_1_4]
else:
sdf_renderer = SDFRenderer(decoder,
camera.intrinsic,
img_hw=img_hw,
march_step=100,
buffer_size=args.buffer_size,
threshold=args.threshold,
ray_marching_ratio=args.ratio,
use_depth2normal=args.use_depth2normal)
renderer_list = [sdf_renderer]
extrinsic = torch.from_numpy(camera.extrinsic).float().cuda()
if args.oracle:
num_iters = 1
else:
num_iters = args.num_iters
# optimization start
latent_tensor, optimizer_latent = optimize_single_view(
renderer_list,
evaluator,
optimizer_latent,
latent_tensor,
extrinsic,
gt_pack,
weight_dict,
optimizer_type="shape",
num_iters=num_iters,
points_gt=points_gt,
test_step=args.test_step,
profile=args.profile,
visualizer=visualizer,
ray_marching_type=args.method,
vis_folder=args.vis_folder)
def main():
import torch.multiprocessing as mp
torch.autograd.set_detect_anomaly(True)
mp.set_start_method('spawn')
import argparse
arg_parser = argparse.ArgumentParser(
description="Color training pipeline."
)
arg_parser.add_argument('-g', '--gpu', default='0', help='gpu id.')
arg_parser.add_argument("--checkpoint", "-c", dest="checkpoint", default="2000",
help='The checkpoint weights to use. This can be a number indicated an epoch or "latest" '
+ "for the latest weights (this is the default)",
)
arg_parser.add_argument('--test_step', '-t', type=int, default=1, help='test step.')
arg_parser.add_argument('--data_path', default='data/demo_multiview_real/', help='path to the dataset.') #TODO need to change the default
arg_parser.add_argument('--obj_name', default='chairs', help='deepsdf class model for experiments. (currently only support "chairs" ')
arg_parser.add_argument('--visualize', action='store_true', help='visualization flag.')
arg_parser.add_argument('--scale', type=float, default=0.2, help='scale the size of input image.')
args = arg_parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
################################
num_views_per_round = 6
num_epoch = 5
sep_dist = 1
refine_sim = True
ini_mean_shape = False
################################
# load data
exp_dir = os.path.join('deepsdf/experiments/', args.obj_name)
out_dir = 'vis/demo_multiview_real'
upper_loader = LoaderMultiReal(args.data_path, scale=args.scale, refine_sim=refine_sim)
for instance_num in range(len(upper_loader)):
# load data
instance_name, imgs, _, cameras, sim_mtrx_ini = upper_loader[instance_num]
vis_dir = os.path.join(out_dir, '{}'.format(instance_num))
os.makedirs(vis_dir, exist_ok=True)
total_num_img = len(imgs)
# RANDOMLY initialize shape code
latent_size = 256
std_ = 0.1
if ini_mean_shape:
shape_code = torch.zeros(1, latent_size)
else:
shape_code = torch.ones(1, latent_size).normal_(mean=0, std=std_)
shape_code = shape_code.float().cuda()
shape_code.requires_grad = True
if refine_sim:
sim3 = easydict.EasyDict({
"scale": torch.tensor(0.,requires_grad=True, device="cuda"),
"rot": torch.tensor([0.,0.,0.],requires_grad=True, device="cuda"),
"trans": torch.tensor([0.,0.,0.],requires_grad=True, device="cuda"),
})
optim_list = [{ "params": [v for k,v in sim3.items()], "lr": LR },
{ "params": [shape_code], "lr": LR }]
optimizer_latent = torch.optim.Adam(optim_list)
else:
optimizer_latent = torch.optim.Adam([shape_code] , lr=LR)
sim3 = None
decoder = load_decoder(exp_dir, args.checkpoint)
decoder = decoder.module.cuda()
img_h, img_w = imgs[0].shape[0], imgs[0].shape[1]
img_hw = (img_h, img_w)
print('Image size: {0}.'. format(img_hw))
sdf_renderer = SDFRenderer_warp(decoder, cameras[0].intrinsic, march_step=200, buffer_size=1, threshold=THRESHOLD, transform_matrix=np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))
evaluator = Evaluator(decoder)
visualizer = Visualizer(img_hw)
weight_list = {}
weight_list['color'] = 5.0
weight_list['l2reg'] = 1.0
shape_code, optimizer_latent = optimize_multi_view(sdf_renderer, evaluator, shape_code, optimizer_latent, imgs, cameras, weight_list, num_views_per_round=num_views_per_round, num_iters=num_epoch, sep_dist=sep_dist, test_step=args.test_step, visualizer=visualizer, points_gt=None, sim3=sim3, sim3_init=sim_mtrx_ini, vis_flag=args.visualize, vis_dir=vis_dir)
print('Finished. check results {}'.format(out_dir))
def train(args):
# initialize output_dir
if not os.path.exists(args.vis_folder):
os.makedirs(args.vis_folder)
#########################################################################
# load data
#########################################################################
mesh_data_dir, experiment_directory, split_file, synthetic_data_dir = init_info(
)
upper_loader = LoaderSingle(synthetic_data_dir, mesh_data_dir,
experiment_directory, split_file)
shape_md5, image_data, mesh_data, camera, depth = upper_loader[6]
img, _, normal, _ = image_data
gt_samples, norm_params = mesh_data
points_gt = np.array(gt_samples.vertices)
points_gt = (points_gt + norm_params['offset']) * norm_params['scale']
gt_pack = {}
gt_pack['depth'] = torch.from_numpy(depth).cuda()
gt_pack['normal'] = torch.from_numpy(normal).cuda()
gt_pack['silhouette'] = torch.from_numpy(depth < 1e5).type(
torch.uint8).cuda()
# # visualize gt
# cv2.imwrite('img.png', img)
# visualize_depth('test0.png', depth)
# with open('camera.pkl', 'wb') as f:
# pickle.dump(camera, f)
#########################################################################
# prepare tensor
#########################################################################
decoder = load_decoder(experiment_directory, args.checkpoint)
decoder = decoder.module.cuda()
evaluator = Evaluator(decoder)
latent_code_dir = os.path.join(synthetic_data_dir, 'latent_codes',
'{0}.pth'.format(shape_md5))
latent_code = torch.load(latent_code_dir)
latent_tensor = latent_code[0].detach().cpu()
latent_size = latent_code.shape[-1]
latent_tensor = latent_tensor.float().cuda()
latent_tensor.requires_grad = False
K, RT = camera.intrinsic, camera.extrinsic
camera_tensor = get_tensor_from_camera(RT)
std_cam = 1e-1
if args.use_gt_camera:
camera_tensor = camera_tensor
else:
camera_tensor = camera_tensor + torch.ones(
camera_tensor.shape[0]).normal_(mean=0, std=std_cam)
from torch.autograd import Variable
camera_tensor = Variable(camera_tensor.cuda(), requires_grad=True)
optimizer_camera = torch.optim.Adam([camera_tensor], lr=args.lr)
#########################################################################
# optimization
#########################################################################
weight_dict = {}
weight_dict['w_depth'] = 10.0
weight_dict['w_normal'] = 10.0
weight_dict['w_mask_gt'] = 1.0
weight_dict['w_mask_out'] = 1.0
weight_dict['w_l2reg'] = 1.0
img_h, img_w = img.shape[0], img.shape[1]
img_hw = (img_h, img_w)
print('Image size: {0}.'.format(img_hw))
if args.visualize:
visualizer = Visualizer(img_hw)
else:
visualizer = None
# initialize renderer
sdf_renderer = SDFRenderer(decoder,
camera.intrinsic,
img_hw=img_hw,
march_step=100,
buffer_size=3,
threshold=args.threshold,
ray_marching_ratio=args.ratio,
use_depth2normal=args.use_depth2normal)
renderer_list = [sdf_renderer]
if args.oracle:
num_iters = 1
else:
num_iters = args.num_iters
# optimization start
camera_tensor, optimizer_latent = optimize_single_view(
renderer_list,
evaluator,
optimizer_camera,
latent_tensor,
camera_tensor,
gt_pack,
weight_dict,
optimizer_type="camera",
num_iters=num_iters,
points_gt=points_gt,
test_step=args.test_step,
profile=args.profile,
visualizer=visualizer,
ray_marching_type=args.method,
vis_folder=args.vis_folder)
def main():
import torch.multiprocessing as mp
mp.set_start_method('spawn')
import argparse
arg_parser = argparse.ArgumentParser(
description="Color training pipeline.")
arg_parser.add_argument('-g', '--gpu', default='0', help='gpu id.')
arg_parser.add_argument(
"--checkpoint",
"-c",
dest="checkpoint",
default="2000",
help=
'The checkpoint weights to use. This can be a number indicated an epoch or "latest" '
+ "for the latest weights (this is the default)",
)
arg_parser.add_argument('--test_step',
'-t',
type=int,
default=5,
help='test step.')
arg_parser.add_argument('--visualize',
action='store_true',
help='visualization flag.')
arg_parser.add_argument('--data_path',
default='data/demo_multiview_syn/',
help='path to PMO dataset.')
arg_parser.add_argument(
'--obj_name',
default='cars',
help=
'deepsdf class model for experiments. (support "planes", "chairs", "cars"'
)
arg_parser.add_argument(
'--scale',
type=float,
default=0.5,
help='scale the size of input image, 224x224 -> 112x112.')
arg_parser.add_argument(
'--focal',
type=float,
default=None,
help='resize the image and change focal length, try 2')
arg_parser.add_argument('--full',
action='store_true',
help='run over all PMO data, otherwise run demo')
args = arg_parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
################################
num_sample_points = 10000
num_views_per_round = 8
################################
# load data
ind = class_type.index(args.obj_name)
class_id = class_ID[ind]
exp_dir = os.path.join('deepsdf/experiments/', args.obj_name)
upper_loader = LoaderMultiPMO(args.data_path,
class_id,
scale=args.scale,
num_points=num_sample_points,
focal=args.focal)
if args.full:
total_num_instance = 50 # consider 50 instances in total
out_dir = os.path.join('vis/multiview_syn/', args.obj_name)
else:
total_num_instance = len(upper_loader) # demo data
out_dir = os.path.join('vis/demo_multiview_syn/', args.obj_name)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
cf_dist1_total = 0.0
cf_dist2_total = 0.0
for instance_num in range(total_num_instance):
vis_dir = os.path.join(out_dir, '{}'.format(instance_num))
if not os.path.exists(vis_dir):
os.makedirs(vis_dir)
instance_name, imgs, masks, cameras, points_gt = upper_loader[
instance_num]
# RANDOMLY initialize shape code
latent_size = 256
std_ = 0.1
shape_code = torch.ones(1, latent_size).normal_(mean=0, std=std_)
shape_code = shape_code.float().cuda()
shape_code.requires_grad = True
decoder = load_decoder(exp_dir, args.checkpoint)
decoder = decoder.module.cuda()
optimizer_latent = torch.optim.Adam([shape_code], lr=LR)
img_h, img_w = imgs[0].shape[0], imgs[0].shape[1]
img_hw = (img_h, img_w)
print('Image size: {0}.'.format(img_hw))
sdf_renderer = SDFRenderer_warp(decoder,
cameras[0].intrinsic,
march_step=100,
buffer_size=1,
threshold=THRESHOLD)
evaluator = Evaluator(decoder)
visualizer = Visualizer(img_hw)
weight_list = {}
weight_list['color'] = 5.0
weight_list['l2reg'] = 1.0
shape_code, optimizer_latent = optimize_multi_view(
sdf_renderer,
evaluator,
shape_code,
optimizer_latent,
imgs,
cameras,
weight_list,
num_views_per_round=num_views_per_round,
num_iters=20,
num_sample_points=num_sample_points,
visualizer=visualizer,
points_gt=points_gt,
vis_dir=vis_dir,
vis_flag=args.visualize,
full_flag=args.full)
if args.full:
# final evaluation
points_tmp = evaluator.latent_vec_to_points(
shape_code, num_points=num_sample_points, silent=True)
dist1, dist2 = evaluator.compute_chamfer_distance(points_gt,
points_tmp,
separate=True)
cf_dist1_total += dist1 * 1000
cf_dist2_total += dist2 * 1000
if args.full:
print('Final Average Chamfer Loss: ',
cf_dist1_total / total_num_instance,
cf_dist2_total / total_num_instance)
print('Finished. check results {}'.format(out_dir))