| 112 | self.sdf_near_surf = self.sdf_near_surf / self.sdf_threshold |
| 113 | |
| 114 | def _sample_batch(self, batch_size): |
| 115 | n_grid = int(batch_size * self.vol_ratio) |
| 116 | n_surf = batch_size - n_grid |
| 117 | |
| 118 | grid_idx = torch.randint(0, self.pts_grid.shape[0], (n_grid,), device=self.device) |
| 119 | surf_idx = torch.randint(0, self.pts_near_surf.shape[0], (n_surf,), device=self.device) |
| 120 | |
| 121 | pts_b = torch.cat([self.pts_grid[grid_idx], self.pts_near_surf[surf_idx]], dim=0) |
| 122 | sdf_b = torch.cat([self.sdf_grid[grid_idx], self.sdf_near_surf[surf_idx]], dim=0) |
| 123 | if self.data_type != "sdf": |
| 124 | tex_b = torch.cat([self.tex_grid[grid_idx], self.tex_near_surf[surf_idx]], dim=0) |
| 125 | return {"pts": pts_b, "sdf": sdf_b, "tex": tex_b} |
| 126 | else: |
| 127 | return {"pts": pts_b, "sdf": sdf_b} |
| 128 | |
| 129 | def _set_optimizer(self, lr, min_lr_ratio=0.01): |
| 130 | """set optimizer and lr scheduler used in training""" |