Get permutation from pointcloud to input voxel coords.
(input_coords, pointcloud, transformation, label_map,
voxel_output, voxel_pred)
| 669 | o3d.visualization.draw_geometries(([*vis_pcd])) |
| 670 | |
| 671 | def permute_pointcloud(input_coords, pointcloud, transformation, label_map, |
| 672 | voxel_output, voxel_pred): |
| 673 | """Get permutation from pointcloud to input voxel coords.""" |
| 674 | |
| 675 | def _hash_coords(coords, coords_min, coords_dim): |
| 676 | return np.ravel_multi_index((coords - coords_min).T, coords_dim) |
| 677 | |
| 678 | # Validate input. |
| 679 | input_batch_size = input_coords[:, -1].max().item() |
| 680 | pointcloud_batch_size = pointcloud[:, -1].max().int().item() |
| 681 | transformation_batch_size = transformation[:, -1].max().int().item() |
| 682 | assert input_batch_size == pointcloud_batch_size == transformation_batch_size |
| 683 | pointcloud_permutation, pointcloud_target = [], [] |
| 684 | |
| 685 | # Process each batch. |
| 686 | for i in range(input_batch_size + 1): |
| 687 | # Filter batch from the data. |
| 688 | input_coords_mask_b = input_coords[:, -1] == i |
| 689 | input_coords_b = (input_coords[input_coords_mask_b])[:, :-1].numpy() |
| 690 | pointcloud_b = pointcloud[pointcloud[:, -1] == i, :-1].numpy() |
| 691 | transformation_b = transformation[i, :-1].reshape(4, 4).numpy() |
| 692 | # Transform original pointcloud to voxel space. |
| 693 | original_coords1 = np.hstack((pointcloud_b[:, :3], np.ones((pointcloud_b.shape[0], 1)))) |
| 694 | original_vcoords = np.floor(original_coords1 @ transformation_b.T)[:, :3].astype(int) |
| 695 | # Hash input and voxel coordinates to flat coordinate. |
| 696 | vcoords_all = np.vstack((input_coords_b, original_vcoords)) |
| 697 | vcoords_min = vcoords_all.min(0) |
| 698 | vcoords_dims = vcoords_all.max(0) - vcoords_all.min(0) + 1 |
| 699 | input_coords_key = _hash_coords(input_coords_b, vcoords_min, vcoords_dims) |
| 700 | original_vcoords_key = _hash_coords(original_vcoords, vcoords_min, vcoords_dims) |
| 701 | # Query voxel predictions from original pointcloud. |
| 702 | key_to_idx = dict(zip(input_coords_key, range(len(input_coords_key)))) |
| 703 | pointcloud_permutation.append( |
| 704 | np.array([key_to_idx.get(i, -1) for i in original_vcoords_key])) |
| 705 | pointcloud_target.append(pointcloud_b[:, -1].astype(int)) |
| 706 | pointcloud_permutation = np.concatenate(pointcloud_permutation) |
| 707 | # Prepare pointcloud permutation array. |
| 708 | pointcloud_permutation = torch.from_numpy(pointcloud_permutation) |
| 709 | permutation_mask = pointcloud_permutation >= 0 |
| 710 | permutation_valid = pointcloud_permutation[permutation_mask] |
| 711 | # Permute voxel output to pointcloud. |
| 712 | pointcloud_output = torch.zeros(pointcloud.shape[0], voxel_output.shape[1]).to(voxel_output) |
| 713 | pointcloud_output[permutation_mask] = voxel_output[permutation_valid] |
| 714 | # Permute voxel prediction to pointcloud. |
| 715 | # NOTE: Invalid points (points found in pointcloud but not in the voxel) are mapped to 0. |
| 716 | pointcloud_pred = torch.ones(pointcloud.shape[0]).int().to(voxel_pred) * 0 |
| 717 | pointcloud_pred[permutation_mask] = voxel_pred[permutation_valid] |
| 718 | # Map pointcloud target to respect dataset IGNORE_LABELS |
| 719 | pointcloud_target = torch.from_numpy( |
| 720 | np.array([label_map[i] for i in np.concatenate(pointcloud_target)])).int() |
| 721 | return pointcloud_output, pointcloud_pred, pointcloud_target |
| 722 | |
| 723 | |
| 724 | def matrix_nms(cate_labels, seg_masks, sum_masks, cate_scores, sigma=2.0, kernel='gaussian', eps=10e-9, nms_thr=0.5): |
nothing calls this directly
no test coverage detected