(
self,
output_dir: str,
overwrite: bool = False,
save_png: bool = True,
save_pt: bool = True,
save_gif: bool = True,
save_video: bool = True,
gif_fps: float = 10,
background_color: T.Union[float, T.List[float]] = 1.,
global_min_depth: float = None,
global_max_depth: float = None,
hit_only: bool = True,
)
| 2464 | self.camera.load_state_dict(state_dict.get('camera', None)) |
| 2465 | |
| 2466 | def save( |
| 2467 | self, |
| 2468 | output_dir: str, |
| 2469 | overwrite: bool = False, |
| 2470 | save_png: bool = True, |
| 2471 | save_pt: bool = True, |
| 2472 | save_gif: bool = True, |
| 2473 | save_video: bool = True, |
| 2474 | gif_fps: float = 10, |
| 2475 | background_color: T.Union[float, T.List[float]] = 1., |
| 2476 | global_min_depth: float = None, |
| 2477 | global_max_depth: float = None, |
| 2478 | hit_only: bool = True, |
| 2479 | ): |
| 2480 | |
| 2481 | if isinstance(background_color, (int, float)): |
| 2482 | background_color = [background_color] * 3 |
| 2483 | |
| 2484 | if os.path.exists(output_dir) and not overwrite: |
| 2485 | raise RuntimeError(f'{output_dir} exists') |
| 2486 | os.makedirs(output_dir, exist_ok=True) |
| 2487 | |
| 2488 | if save_pt: |
| 2489 | filename = os.path.join(output_dir, 'state_dict.pt') |
| 2490 | torch.save(self.state_dict(), filename) |
| 2491 | |
| 2492 | # deal with hit_map |
| 2493 | b, q, h, w, _3 = self.rgb.shape |
| 2494 | if hit_only and self.hit_map is not None: |
| 2495 | hit_map = (self.hit_map > 0.5).float() # (b, q, h, w) |
| 2496 | else: |
| 2497 | hit_map = torch.ones(b, q, h, w, dtype=self.rgb.dtype, device=self.rgb.device) # (b, q, h, w) |
| 2498 | |
| 2499 | background_img = torch.ones_like(self.rgb) # (b, q, h, w, 3) |
| 2500 | for c in range(3): |
| 2501 | background_img[..., c] = background_color[c] |
| 2502 | |
| 2503 | # deal with depth |
| 2504 | if self.depth is not None: |
| 2505 | masked_depth = self.depth * hit_map |
| 2506 | else: |
| 2507 | masked_depth = None |
| 2508 | masked_rgb = self.rgb * hit_map.unsqueeze(-1) |
| 2509 | masked_rgb = masked_rgb + (1 - hit_map).unsqueeze(-1) * background_img |
| 2510 | if self.normal_w is not None: |
| 2511 | masked_normal_w = self.normal_w * hit_map.unsqueeze(-1) |
| 2512 | masked_normal_w = masked_normal_w + (1 - hit_map).unsqueeze(-1) * background_img |
| 2513 | else: |
| 2514 | masked_normal_w = None |
| 2515 | |
| 2516 | if save_png: |
| 2517 | # normal: [-1, 1] -> [0, 1] |
| 2518 | if self.normal_w is not None: |
| 2519 | normal_w = (masked_normal_w + 1) / 2. |
| 2520 | else: |
| 2521 | normal_w = None |
| 2522 | |
| 2523 | for ib in range(self.rgb.size(0)): |
no test coverage detected