| 14 | |
| 15 | |
| 16 | class BasicLogger: |
| 17 | def __init__(self, args, for_eva=False) -> None: |
| 18 | self.args = args |
| 19 | self.log_dir = osp.join( |
| 20 | args.log_dir, args.exp_name, self.get_random_time_str()) |
| 21 | self.img_dir = osp.join(self.log_dir, "imgs") |
| 22 | self.mesh_dir = osp.join(self.log_dir, "mesh") |
| 23 | self.ckpt_dir = osp.join(self.log_dir, "ckpt") |
| 24 | self.backup_dir = osp.join(self.log_dir, "bak") |
| 25 | self.misc_dir = osp.join(self.log_dir, "misc") |
| 26 | self.for_eva = for_eva |
| 27 | if not for_eva: |
| 28 | os.makedirs(self.img_dir) |
| 29 | os.makedirs(self.ckpt_dir) |
| 30 | os.makedirs(self.mesh_dir) |
| 31 | os.makedirs(self.misc_dir) |
| 32 | os.makedirs(self.backup_dir) |
| 33 | |
| 34 | self.log_config(args) |
| 35 | |
| 36 | def get_random_time_str(self): |
| 37 | return datetime.strftime(datetime.now(), "%Y-%m-%d-%H-%M-%S") |
| 38 | |
| 39 | def log_ckpt(self, mapper, name): |
| 40 | decoder_state = {f: v.cpu() |
| 41 | for f, v in mapper.decoder.state_dict().items()} |
| 42 | map_state = {f: v.cpu() for f, v in mapper.map_states.items()} |
| 43 | sdf_priors = mapper.sdf_priors.cpu() |
| 44 | torch.save({ |
| 45 | "decoder_state": decoder_state, |
| 46 | "map_state": map_state, |
| 47 | "sdf_priors": sdf_priors}, |
| 48 | os.path.join(self.ckpt_dir, name)) |
| 49 | |
| 50 | def log_config(self, config): |
| 51 | out_path = osp.join(self.backup_dir, "config.yaml") |
| 52 | yaml.dump(vars(config), open(out_path, 'w')) |
| 53 | |
| 54 | def log_mesh(self, mesh, name="final_mesh.ply"): |
| 55 | out_path = osp.join(self.mesh_dir, name) |
| 56 | o3d.io.write_triangle_mesh(out_path, mesh) |
| 57 | |
| 58 | def log_point_cloud(self, pcd, name="final_points.ply"): |
| 59 | out_path = osp.join(self.mesh_dir, name) |
| 60 | o3d.io.write_point_cloud(out_path, pcd) |
| 61 | |
| 62 | def log_numpy_data(self, data, name, ind=None): |
| 63 | if isinstance(data, torch.Tensor): |
| 64 | data = data.detach().cpu().numpy() |
| 65 | if ind is not None: |
| 66 | np.save(osp.join(self.misc_dir, "{}-{:05d}.npy".format(name, ind)), data) |
| 67 | else: |
| 68 | np.save(osp.join(self.misc_dir, f"{name}.npy"), data) |
| 69 | |
| 70 | def log_debug_data(self, data, idx): |
| 71 | with open(os.path.join(self.misc_dir, f"scene_data_{idx}.pkl"), 'wb') as f: |
| 72 | pickle.dump(data, f) |
| 73 |
no outgoing calls
no test coverage detected