| 157 | ThreadPoolExecutor(max_workers=32) as saver_executor: |
| 158 | |
| 159 | def loader(task): |
| 160 | sha256, view_idx = task |
| 161 | try: |
| 162 | # Check if output file already exists, skip if so |
| 163 | output_path = os.path.join( |
| 164 | opt.shape_latent_root, |
| 165 | 'shape_latents', |
| 166 | latent_view_name, |
| 167 | sha256, |
| 168 | f'view{view_idx:02d}.npz' |
| 169 | ) |
| 170 | if os.path.exists(output_path): |
| 171 | load_queue.put((sha256, view_idx, None, None)) |
| 172 | return |
| 173 | |
| 174 | # dual_grid_view path: dual_grid_view_{res}/{sha256}/view{idx:02d}.vxz |
| 175 | vxz_path = os.path.join( |
| 176 | opt.dual_grid_root, |
| 177 | f'dual_grid_view_{opt.resolution}', |
| 178 | sha256, |
| 179 | f'view{view_idx:02d}.vxz' |
| 180 | ) |
| 181 | |
| 182 | if not os.path.exists(vxz_path): |
| 183 | print(f"[Loader Skip] {sha256}/view{view_idx:02d}: vxz file not found") |
| 184 | load_queue.put((sha256, view_idx, None, None)) |
| 185 | return |
| 186 | |
| 187 | coords, attr = o_voxel.io.read_vxz(vxz_path, num_threads=4) |
| 188 | vertices = sp.SparseTensor( |
| 189 | (attr['vertices'] / 255.0).float(), |
| 190 | torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1), |
| 191 | ) |
| 192 | intersected = vertices.replace(torch.cat([ |
| 193 | attr['intersected'] % 2, |
| 194 | attr['intersected'] // 2 % 2, |
| 195 | attr['intersected'] // 4 % 2, |
| 196 | ], dim=-1).bool()) |
| 197 | load_queue.put((sha256, view_idx, vertices, intersected)) |
| 198 | except Exception as e: |
| 199 | print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}") |
| 200 | load_queue.put((sha256, view_idx, None, None)) |
| 201 | |
| 202 | loader_executor.map(loader, tasks) |
| 203 | |