(
self,
output_dir: str,
overwrite: bool = False,
save_ply: bool = True,
save_pt: bool = True,
cylinder_radius: float = 0.1,
cone_radius: float = None,
max_cone_height: float = None,
end_xyz_w: torch.Tensor = None, # same shape as origins_w
)
| 1135 | setattr(self, name, state_dict.get(name, None)) |
| 1136 | |
| 1137 | def save( |
| 1138 | self, |
| 1139 | output_dir: str, |
| 1140 | overwrite: bool = False, |
| 1141 | save_ply: bool = True, |
| 1142 | save_pt: bool = True, |
| 1143 | cylinder_radius: float = 0.1, |
| 1144 | cone_radius: float = None, |
| 1145 | max_cone_height: float = None, |
| 1146 | end_xyz_w: torch.Tensor = None, # same shape as origins_w |
| 1147 | ): |
| 1148 | if self.origins_w is None or self.directions_w is None: |
| 1149 | return |
| 1150 | |
| 1151 | if cone_radius is None: |
| 1152 | cone_radius = cylinder_radius * 1.5 |
| 1153 | |
| 1154 | if max_cone_height is None: |
| 1155 | max_cone_height = cylinder_radius * 2. |
| 1156 | |
| 1157 | if os.path.exists(output_dir) and not overwrite: |
| 1158 | raise RuntimeError(f'output dir {output_dir} exists') |
| 1159 | os.makedirs(output_dir, exist_ok=True) |
| 1160 | |
| 1161 | if save_pt: |
| 1162 | filename = os.path.join(output_dir, 'state_dict.pt') |
| 1163 | torch.save(self.state_dict(), filename) |
| 1164 | |
| 1165 | if end_xyz_w is None: |
| 1166 | ray_ts = torch.ones(*self.origins_w.shape[:-1], device=self.origins_w.device) |
| 1167 | add_end_ball = False |
| 1168 | else: |
| 1169 | ray_ts = torch.linalg.vector_norm(end_xyz_w - self.origins_w, ord=2, dim=-1) |
| 1170 | add_end_ball = True |
| 1171 | |
| 1172 | b, *m_shape, _3 = self.origins_w.shape |
| 1173 | origins_w = self.origins_w.reshape(b, -1, 3) |
| 1174 | directions_w = self.directions_w.reshape(b, -1, 3) |
| 1175 | ray_ts = ray_ts.reshape(b, -1) |
| 1176 | if save_ply: |
| 1177 | # we are going to save individual camera frames |
| 1178 | for ib in range(origins_w.size(0)): |
| 1179 | sub_dir = os.path.join(output_dir, f'batch_{ib}') |
| 1180 | os.makedirs(sub_dir, exist_ok=True) |
| 1181 | |
| 1182 | for im in range(origins_w.size(1)): |
| 1183 | |
| 1184 | R = rigid_motion.get_min_R( |
| 1185 | v1=np.array([0, 0, 1.], dtype=np.float32), |
| 1186 | v2=directions_w[ib, im].detach().cpu().numpy(), |
| 1187 | ) |
| 1188 | H_c2w = np.eye(4) |
| 1189 | H_c2w[:3, :3] = R |
| 1190 | H_c2w[:3, 3] = origins_w[ib, im].detach().cpu().float().numpy() |
| 1191 | |
| 1192 | # ray origin |
| 1193 | mesh = o3d.geometry.TriangleMesh.create_sphere( |
| 1194 | radius=cylinder_radius * 1.1, |
no test coverage detected