def between_levels(self):
self.logger.start()
self.update_target()
if self.source_mesh is None:
last_mesh_path = f'{self.opt.cp_folder}/{self.opt.mesh_name}_level{self.level - 1 + self.refinement:02d}.obj'
if self.level > 0 and os.path.isfile(last_mesh_path):
self.source_mesh = mesh_utils.to(mesh_utils.load_mesh(last_mesh_path), self.device)
self.source_mesh = self.source_mesh[0] * (2 ** self.level), self.source_mesh[1]
self.target_mesh = self.target_mesh[0] * (2 ** self.level), self.target_mesh[1]
else:
mesh_path = f'{const.RAW_MESHES}/{self.opt.mesh_name}_template.obj'
if not os.path.isfile(mesh_path) or not self.opt.pre_template:
# mesh_path = f'{const.RAW_MESHES}/icosahedron.obj'
self.source_mesh = mesh_utils.load_real_mesh(self.opt.template_name, self.opt.template_start, False)
else:
self.source_mesh = mesh_utils.load_mesh(mesh_path)
self.source_mesh = mesh_utils.to(mesh_utils.to_unit_cube(self.source_mesh)[0], self.device)
scale = 1 / mesh_utils.edge_lengths(self.source_mesh).mean().item()
self.target_mesh = self.target_mesh[0] * scale, self.target_mesh[1]
self.source_mesh = self.source_mesh[0] * scale, self.source_mesh[1]
self.target_transform = self.target_transform[0], self.target_transform[1] * scale
if self.opt.template_start > 0 or os.path.isfile(mesh_path):
mesh_utils.export_mesh(self.source_mesh, f'{self.opt.cp_folder}/{self.opt.mesh_name}_template.obj')
if self.level != 0 and not self.refinement:
self.source_mesh = mesh_utils.Upsampler(self.source_mesh)(self.source_mesh)
self.ds_source = mesh_utils.VerticesDS(self.source_mesh).to(self.device)
self.save_scales()
def before_level(self) -> [Optimizer, Optimizer, MeshHandler]:
self.logger.start()
self.generator.dup(self.level)
self.discriminator.dup(self.level)
self.optimizer_generator = factory.OptimizerLC(self.opt, 'generator', self.generator.get_level(self.level))
self.optimizer_discriminator = factory.OptimizerLC(self.opt, 'discriminator', self.discriminator.get_level(self.level))
self.real_mesh = MeshHandler(mesh_utils.load_real_mesh(self.opt.mesh_name,self.opt.start_level + self.level),
self.opt, self.opt.start_level + self.level).to(self.device)
if self.opt.inside_out:
self.real_mesh_flipped = self.real_mesh.copy().flip()
def load_template_mesh(opt: Options, level) -> Tuple[str, T_Mesh]:
mesh_path = f'{DATA_ROOT}/{opt.mesh_name}/{opt.mesh_name}_template.obj'
if not os.path.isfile(mesh_path):
return opt.template_name, mesh_utils.load_real_mesh(opt.template_name, level)
else:
mesh = mesh_utils.scale_mesh(mesh_path, False, opt.mesh_name, 0)
if level > 0:
mesh_handler = MeshHandler(mesh, opt, 0)
for i in range(level):
mesh_handler.upsample()
mesh = mesh_handler.raw_mesh
return f'{opt.mesh_name}_template', mesh
def __call__(self,
mesh: Union[str, MeshHandler, N],
start: int,
end: int,
zero_places: NoiseT = ()) -> MeshHandler:
start, end = self.trim(start, end)
if len(zero_places) == 1:
zero_places = zero_places * (end - start + 1)
if mesh is None:
mesh = MeshHandler(
mesh_utils.load_real_mesh(self.opt.template_name, start),
self.opt, 0).to(self.device)
elif type(mesh) is str:
mesh = MeshHandler(mesh, self.opt, 0).to(self.device)
z = self.get_z_sequence(mesh, end - start)
for i in range(min(len(z), len(zero_places))):
if zero_places[i]:
z[i] = 0
remeshed = self.generator.forward(mesh, z, end, start, upsample=True)
return remeshed
def synthesize(args):
device = CUDA(0)
# Generating Training Data
gt_paths = [
f'{const.DATA_ROOT}/{args.train_mesh}/{args.train_mesh}_level{i:02d}.obj'
for i in range(6)
]
is_generated = all(list(os.path.isfile(gt_path) for gt_path in gt_paths))
if (not is_generated) or args.no_cache:
gen_args = options.GtOptions(tag='demo',
mesh_name=args.train_mesh,
template_name='sphere',
num_levels=6)
gt_gen = GroundTruthGenerator(gen_args, device)
print("Finished generating training data with " + args.train_mesh,
flush=True)
# Training Synthesizer
options_path = f'{const.PROJECT_ROOT}/checkpoints/{args.train_mesh}_demo/options.pkl'
models_path = f'{const.PROJECT_ROOT}/checkpoints/{args.train_mesh}_demo/SingleMeshGenerator.pth'
is_trained = os.path.isfile(options_path) and os.path.isfile(models_path)
train_args = options.TrainOption(tag='demo',
mesh_name=args.train_mesh,
template_name='sphere',
num_levels=6)
if (not is_trained) or args.no_cache:
trainer = Trainer(train_args, device)
trainer.train()
print("Finished training with " + args.train_mesh, flush=True)
# Synthesizing Input
m2m = Mesh2Mesh(train_args, CPU)
mesh = mesh_utils.load_real_mesh(args.input_mesh, 0, True)
out = m2m(mesh, 2, 5, 0)
out.export(f'{const.RAW_MESHES}/{args.input_mesh}_hi')
print("Finished synthesizing input on " + args.input_mesh, flush=True)
f'gen {self.opt.mesh_name} {i * num_samples + j +1:02d} / {len(self.generator.levels) * num_samples}'
)
def __call__(self, start_level: int):
with torch.no_grad():
if start_level < 0:
start_level = len(self)
start_level = min(len(self), start_level)
z = self.compose_z(start_level)
return self.generator(self.template.copy(), z, len(self) - 1)
if __name__ == '__main__':
opt_ = options.Options()
opt_.parse_cmdline()
device = CPU
with_noise = False
if opt_.gen_mode == 'generate':
mg = MeshGen(opt_, device)
mg.generate_all(opt_.num_gen_samples)
elif opt_.gen_mode == 'animate':
m2m = Mesh2Mesh(opt_, device)
in_mesh = MeshInference(
opt_.target, mesh_utils.load_real_mesh(opt_.target, 0, True), opt_,
0).to(device)
m2m.animate(in_mesh,
opt_.gen_levels[0],
opt_.gen_levels[1],
0, (12, 17),
zero_places=(0, 0, 1, 1, 1))