(data_path, device="cuda")
| 21 | PART_NORMALIZE_SCALE = 0.7 |
| 22 | |
| 23 | def prepare_data(data_path, device="cuda"): |
| 24 | if data_path.endswith(".glb"): |
| 25 | parts_mesh = trimesh.load(data_path) |
| 26 | part_name_list = [] |
| 27 | part_pcd_list = [] |
| 28 | whole_cond_list = [] |
| 29 | part_cond_list = [] |
| 30 | part_local_cond_list = [] |
| 31 | part_center_list = [] |
| 32 | part_scale_list = [] |
| 33 | for i, (name, part_mesh) in enumerate(parts_mesh.geometry.items()): |
| 34 | part_surface_points, face_idx = part_mesh.sample(NUM_SURFACE_SAMPLES, return_index=True) |
| 35 | part_surface_normals = part_mesh.face_normals[face_idx] |
| 36 | part_pcd = np.concatenate([part_surface_points, np.ones_like(part_surface_points[:, :1])*i], axis=-1) |
| 37 | part_pcd_list.append(part_pcd) |
| 38 | |
| 39 | part_surface_points = torch.FloatTensor(part_surface_points) |
| 40 | part_surface_normals = torch.FloatTensor(part_surface_normals) |
| 41 | part_cond = torch.cat([part_surface_points, part_surface_normals], dim=-1) |
| 42 | part_local_cond = part_cond.clone() |
| 43 | part_cond_max = part_local_cond[:, :3].max(dim=0)[0] |
| 44 | part_cond_min = part_local_cond[:, :3].min(dim=0)[0] |
| 45 | part_center_new = (part_cond_max + part_cond_min) / 2 |
| 46 | part_local_cond[:, :3] = part_local_cond[:, :3] - part_center_new |
| 47 | part_scale_new = (part_local_cond[:, :3].abs().max() / (0.95 * PART_NORMALIZE_SCALE)).item() |
| 48 | part_local_cond[:, :3] = part_local_cond[:, :3] / part_scale_new |
| 49 | part_cond_list.append(part_cond) |
| 50 | part_local_cond_list.append(part_local_cond) |
| 51 | part_name_list.append(name) |
| 52 | part_center_list.append(part_center_new) |
| 53 | part_scale_list.append(part_scale_new) |
| 54 | |
| 55 | part_pcd = np.concatenate(part_pcd_list, axis=0) |
| 56 | part_pcd = torch.FloatTensor(part_pcd).to(device) |
| 57 | whole_mesh = parts_mesh.dump(concatenate=True) |
| 58 | whole_surface_points, face_idx = whole_mesh.sample(NUM_SURFACE_SAMPLES, return_index=True) |
| 59 | whole_surface_normals = whole_mesh.face_normals[face_idx] |
| 60 | whole_surface_points = torch.FloatTensor(whole_surface_points) |
| 61 | whole_surface_normals = torch.FloatTensor(whole_surface_normals) |
| 62 | whole_surface_points_tensor = whole_surface_points.to(device) |
| 63 | nearest_idx = nearest(whole_surface_points_tensor, part_pcd[:, :3]) |
| 64 | nearest_part = part_pcd[nearest_idx] |
| 65 | nearest_part = nearest_part[:, 3].cpu() |
| 66 | for i in range(len(part_cond_list)): |
| 67 | surface_points_part_mask = (nearest_part == i).float() |
| 68 | whole_cond = torch.cat([whole_surface_points, whole_surface_normals, surface_points_part_mask[..., None]], dim=-1) |
| 69 | whole_cond_list.append(whole_cond) |
| 70 | |
| 71 | batch_data = { |
| 72 | "whole_cond": torch.stack(whole_cond_list, dim=0).to(device), |
| 73 | "part_cond": torch.stack(part_cond_list, dim=0).to(device), |
| 74 | "part_local_cond": torch.stack(part_local_cond_list, dim=0).to(device), |
| 75 | "part_id_list": part_name_list, |
| 76 | "part_center_list": part_center_list, |
| 77 | "part_scale_list": part_scale_list, |
| 78 | } |
| 79 | else: |
| 80 | raise ValueError("Unsupported file format. Please provide a .glb file.") |
no test coverage detected