| 49 | self.sdf_threshold = None |
| 50 | |
| 51 | def _load_data(self, path, sdf_renorm=False): |
| 52 | data = np.load(path) |
| 53 | self.aabb = torch.from_numpy(data["aabb"]).float().to(self.device) |
| 54 | self.sdf_threshold = float(data["threshold"]) |
| 55 | self.Ka = data["Ka"].tolist() if "Ka" in data else [0, 0, 0] |
| 56 | self.Kd = data["Kd"].tolist() if "Kd" in data else [1, 1, 1] |
| 57 | self.Ks = data["Ks"].tolist() if "Ks" in data else [0.4, 0.4, 0.4] |
| 58 | self.Ns = data["Ns"].tolist() if "Ns" in data else 10 |
| 59 | print("aabb: ", self.aabb) |
| 60 | print("using sdf_threshold: ", self.sdf_threshold) |
| 61 | |
| 62 | pts_grid = data["pts_grid"] |
| 63 | sdf_gird = data["sdf_grid"] |
| 64 | pts_near_surf = data["pts_near_surf"] |
| 65 | sdf_near_surf = data["sdf_near_surf"] |
| 66 | |
| 67 | if self.data_type != "sdf": |
| 68 | tex_grid = data["tex_grid"] |
| 69 | pts_on_surf = data["pts_on_surf"] |
| 70 | tex_on_surf = data["tex_on_surf"] |
| 71 | tex_near_surf = data["tex_near_surf"] |
| 72 | print("pts_grid shape: ", pts_grid.shape) |
| 73 | print("pts_near_surf shape: ", pts_near_surf.shape) |
| 74 | |
| 75 | self.featmap_size = (torch.tensor(pts_grid.shape[:3]).float() * (self.fm_reso / max(pts_grid.shape[:3]))).long().tolist() |
| 76 | self.featmap_size = [int(x // 2 * 2) for x in self.featmap_size] |
| 77 | print("featmap size: ", self.featmap_size) |
| 78 | |
| 79 | # index volume |
| 80 | if self.data_type != "sdf": |
| 81 | input_grid = np.concatenate([sdf_gird[np.newaxis], tex_grid.transpose(3, 0, 1, 2)], axis=0) |
| 82 | else: |
| 83 | input_grid = sdf_gird[np.newaxis] |
| 84 | input_grid = torch.from_numpy(input_grid).float().to(self.device) |
| 85 | required_shape = [x * 2 for x in self.featmap_size] |
| 86 | if input_grid.shape[1] != required_shape[0] or input_grid.shape[2] != required_shape[1] or input_grid.shape[3] != required_shape[2]: |
| 87 | print("resize input_grid from ", input_grid.shape, " to ", required_shape) |
| 88 | input_grid = F.interpolate(input_grid.unsqueeze(0), size=required_shape, mode="trilinear", align_corners=False).squeeze(0) |
| 89 | self.input_grid = input_grid.unsqueeze(0) # [1, C, H, W, D] |
| 90 | print("input grid shape: ", self.input_grid.shape) |
| 91 | |
| 92 | self.pts_grid = torch.from_numpy(pts_grid).float().to(self.device).view(-1, 3) |
| 93 | self.sdf_grid = torch.from_numpy(sdf_gird).float().to(self.device).view(-1, 1).clamp_(-self.sdf_threshold, self.sdf_threshold) |
| 94 | self.pts_near_surf = torch.from_numpy(pts_near_surf).float().to(self.device).view(-1, 3) |
| 95 | self.sdf_near_surf = torch.from_numpy(sdf_near_surf).float().to(self.device).view(-1, 1).clamp_(-self.sdf_threshold, self.sdf_threshold) |
| 96 | |
| 97 | if self.data_type != "sdf": |
| 98 | tex_channels = tex_grid.shape[-1] |
| 99 | self.tex_grid = torch.from_numpy(tex_grid).float().to(self.device).view(-1, tex_channels) |
| 100 | self.pts_on_surf = torch.from_numpy(pts_on_surf).float().to(self.device).view(-1, 3) |
| 101 | self.tex_on_surf = torch.from_numpy(tex_on_surf).float().to(self.device).view(-1, tex_channels) |
| 102 | self.tex_near_surf = torch.from_numpy(tex_near_surf).float().to(self.device).view(-1, tex_channels) |
| 103 | |
| 104 | if self.pts_on_surf.shape[0] > 2_000_000: |
| 105 | print("downsample pts_on_surf from ", self.pts_on_surf.shape[0], " to 2,000,000") |
| 106 | idx = torch.randperm(self.pts_on_surf.shape[0])[:2_000_000] |
| 107 | self.pts_on_surf = self.pts_on_surf[idx] |
| 108 | self.tex_on_surf = self.tex_on_surf[idx] |