def save_optimizer(optimizer_):
nonlocal already_init, optimizer_path, scheduler
if not already_init:
init_folders(optimizer_path)
already_init = True
torch.save(optimizer_.state_dict(), optimizer_path)
def __init__(self, opt: options.Options, cache_length: int):
super(MeshDataset, self).__init__()
self.opt = opt
files_utils.init_folders(self.data_path(0))
self.data_paths = self.get_taxonomy_models_paths()
self.cache_length = cache_length
self.all_data: List[Union[N, mesh_utils.MeshWrap]] = [None] * len(self)
self.delete_cache()
def get_path(prefix: str, saving_index: int) -> str:
nonlocal saving_dict
saving_folder = f'{saving_folders[saving_index]}{prefix}/'
if prefix not in saving_dict[saving_index]:
init_folders(f'{saving_folders[saving_index]}/{prefix}/')
init_prefix(prefix, saving_index, saving_folder)
# saving_dict[saving_index][prefix] += 1
return saving_folder
def plot_mesh(*meshes: Union[T_Mesh, str], save_path: str = '',
ambient_color: T = T((255., 200, 255.)), light_dir: T = T((.5, .5, 1))):
np_images = [render_mesh(mesh, ambient_color, light_dir) for mesh in meshes]
im = make_pretty(np_images)
if save_path:
save_path = files_utils.add_suffix(save_path, '.png')
files_utils.init_folders(save_path)
im.save(save_path)
return im
def generate_ground_truth_meshes(self):
files_utils.init_folders(self.opt.cp_folder + '/')
for i in range(self.opt.start_level, self.opt.start_level + self.opt.num_levels + self.refinement):
self.between_levels()
self.source_mesh = self.optimize()
scaled_mesh = self.adjust_edge_scale()
self.source_mesh = self.source_mesh[0] * 2, self.source_mesh[1]
self.target_mesh = self.target_mesh[0] * 2, self.target_mesh[1]
self.target_transform = self.target_transform[0], self.target_transform[1] * 2
mesh_utils.export_mesh(scaled_mesh, f'{self.opt.cp_folder}/{self.opt.mesh_name}_level{self.level:02d}.obj')
self.logger.stop(False)
self.level += 1
self.save_scales()
def export_mesh(mesh: Union[V_Mesh, T_Mesh], file_name: str):
vs, faces = to(mesh, CPU)
file_name = add_suffix(file_name, '.obj')
init_folders(file_name)
if faces is not None:
faces = faces + 1
with open(file_name, 'w') as f:
for vi, v in enumerate(vs):
f.write("v %f %f %f\n" % (v[0], v[1], v[2]))
if faces is not None:
for face_id in range(faces.shape[0] - 1):
f.write("f %d %d %d\n" % (faces[face_id][0], faces[face_id][1], faces[face_id][2]))
f.write("f %d %d %d" % (faces[-1][0], faces[-1][1], faces[-1][2]))
def save_model(model):
nonlocal already_init, last_recover, model_path
if not already_init:
init_folders(model_path)
already_init = True
if is_model_clean(model):
recover_path = f'{opt.cp_folder}/{model_class_name}_r{last_recover + 1}.pth'
if os.path.isfile(recover_path):
os.remove(recover_path)
torch.save(model.state_dict(), model_path)
shutil.copy(model_path, recover_path)
last_recover = 1 - last_recover
return True
else:
return False
def view(points: V or list,
splits: V or list or tuple = (),
palette=None,
save_path='',
global_transform=None,
titles=(),
rows=1):
points, splits = fix_points(points, splits, global_transform)
bounds = get_bounds(*points)
palette = create_palettes(splits, palette)
np_images = view_backend(points, splits, bounds, palette, titles)
im = make_pretty(np_images, rows=rows)
if save_path:
init_folders(save_path)
im.save(save_path)
elif 'DISPLAY' in os.environ:
im.show()
return im, points
def save_model(model_: SingleModel):
nonlocal already_init, last_recover, model_path
if opt.debug:
return False
if not already_init:
init_folders(model_path)
with open(params_path, 'wb') as f_p:
pickle.dump({'params': opt, 'model_name': model_class_name}, f_p, pickle.HIGHEST_PROTOCOL)
already_init = True
if is_model_clean(model_):
recover_path = f'{opt.cp_folder}/{model_class_name}_recover{last_recover + 1}.pth'
if os.path.isfile(recover_path):
os.remove(recover_path)
torch.save(model_.state_dict(), model_path)
shutil.copy(model_path, recover_path)
last_recover = 1 - last_recover
return True
else:
return False
def save(self):
if not self.already_init:
init_folders(self.optimizer_path)
self.already_init = True
torch.save(self.state_dict(), self.optimizer_path)