(
self,
output_dir: str,
overwrite: bool = False,
save_ply: bool = True,
save_individual_ply: bool = True,
save_pt: bool = True,
world_frame_size: float = 1.,
camera_frame_size: float = 0.1,
scene_meshes: T.Optional[T.List[o3d.geometry.TriangleMesh]] = None,
)
| 2039 | return all_camera_arrows |
| 2040 | |
| 2041 | def save( |
| 2042 | self, |
| 2043 | output_dir: str, |
| 2044 | overwrite: bool = False, |
| 2045 | save_ply: bool = True, |
| 2046 | save_individual_ply: bool = True, |
| 2047 | save_pt: bool = True, |
| 2048 | world_frame_size: float = 1., |
| 2049 | camera_frame_size: float = 0.1, |
| 2050 | scene_meshes: T.Optional[T.List[o3d.geometry.TriangleMesh]] = None, |
| 2051 | ): |
| 2052 | if os.path.exists(output_dir) and not overwrite: |
| 2053 | raise RuntimeError(f'output dir {output_dir} exists') |
| 2054 | os.makedirs(output_dir, exist_ok=True) |
| 2055 | |
| 2056 | if scene_meshes is not None and not isinstance(scene_meshes, (list, tuple)): |
| 2057 | scene_meshes = [scene_meshes] * self.H_c2w.size(0) |
| 2058 | |
| 2059 | if save_pt: |
| 2060 | filename = os.path.join(output_dir, 'state_dict.pt') |
| 2061 | torch.save(self.state_dict(), filename) |
| 2062 | |
| 2063 | if save_individual_ply: |
| 2064 | # save a ply file containing world and camera coordinates |
| 2065 | all_camera_frames = self.get_camera_frames(camera_frame_size=camera_frame_size) |
| 2066 | |
| 2067 | # we are going to save individual camera frames |
| 2068 | for ib in range(self.H_c2w.size(0)): |
| 2069 | sub_dir = os.path.join(output_dir, f'batch_{ib}') |
| 2070 | os.makedirs(sub_dir, exist_ok=True) |
| 2071 | |
| 2072 | for iq in range(self.H_c2w.size(1)): |
| 2073 | filename = os.path.join(sub_dir, f'{iq}.ply') |
| 2074 | o3d.io.write_triangle_mesh( |
| 2075 | filename=filename, |
| 2076 | mesh=all_camera_frames[ib][iq], |
| 2077 | ) |
| 2078 | |
| 2079 | # save world coord |
| 2080 | world_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=world_frame_size) |
| 2081 | filename = os.path.join(sub_dir, f'world.ply') |
| 2082 | o3d.io.write_triangle_mesh( |
| 2083 | filename=filename, |
| 2084 | mesh=world_frame, |
| 2085 | ) |
| 2086 | |
| 2087 | # scene |
| 2088 | if scene_meshes is not None and scene_meshes[ib] is not None: |
| 2089 | filename = os.path.join(sub_dir, f'scene.obj') |
| 2090 | try: |
| 2091 | o3d.io.write_triangle_mesh( |
| 2092 | filename=filename, |
| 2093 | mesh=scene_meshes[ib], |
| 2094 | ) |
| 2095 | except: |
| 2096 | pass |
| 2097 | |
| 2098 | if save_ply: |
nothing calls this directly
no test coverage detected