nbv_reconstruction/utils/pts.py

import numpy as np
import open3d as o3d
import torch

class PtsUtil:

    @staticmethod
    def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005):
        voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
        unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
        return unique_voxels[0]*voxel_size
    
    @staticmethod
    def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
        if point_cloud.shape[0] == 0:
            if require_idx:
                return point_cloud, np.array([])
            return point_cloud
        idx = np.random.choice(len(point_cloud), num_points, replace=True)
        if require_idx:
            return point_cloud[idx], idx
        return point_cloud[idx]
    
    @staticmethod
    def fps_downsample_point_cloud(point_cloud, num_points, require_idx=False):
        N = point_cloud.shape[0]
        mask = np.zeros(N, dtype=bool)

        sampled_indices = np.zeros(num_points, dtype=int)
        sampled_indices[0] = np.random.randint(0, N)
        distances = np.linalg.norm(point_cloud - point_cloud[sampled_indices[0]], axis=1)
        for i in range(1, num_points):
            farthest_index = np.argmax(distances)
            sampled_indices[i] = farthest_index
            mask[farthest_index] = True

            new_distances = np.linalg.norm(point_cloud - point_cloud[farthest_index], axis=1)
            distances = np.minimum(distances, new_distances)
        
        sampled_points = point_cloud[sampled_indices]
        if require_idx:
            return sampled_points, sampled_indices
        return sampled_points
        
    @staticmethod
    def random_downsample_point_cloud_tensor(point_cloud, num_points):
        idx = torch.randint(0, len(point_cloud), (num_points,))
        return point_cloud[idx]
    
    @staticmethod
    def voxelize_points(points, voxel_size):
        voxel_indices = np.floor(points / voxel_size).astype(np.int32)
        unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
        return unique_voxels

    @staticmethod
    def transform_point_cloud(points, pose_mat):
        points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)
        points_h = np.dot(pose_mat, points_h.T).T
        return points_h[:, :3]
    
    @staticmethod
    def get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size=0.005, require_idx=False):
        voxels_L, indices_L = PtsUtil.voxelize_points(point_cloud_L, voxel_size)
        voxels_R, _ = PtsUtil.voxelize_points(point_cloud_R, voxel_size)

        voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
        voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
        overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R)
        mask_L = np.isin(
            indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
        )
        overlapping_points = point_cloud_L[mask_L]
        if require_idx:
            return overlapping_points, mask_L
        return overlapping_points
    
    @staticmethod
    def filter_points(points, normals, cam_pose, theta_limit=45, z_range=(0.2, 0.45)):
        
        """ filter with normal """ 
        normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
        cos_theta = np.dot(normals_normalized, np.array([0, 0, 1]))
        theta = np.arccos(cos_theta) * 180 / np.pi
        idx = theta < theta_limit
        filtered_sampled_points = points[idx]
        filtered_normals = normals[idx]
        
        """ filter with z range """
        points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose))
        idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
        z_filtered_points = filtered_sampled_points[idx]
        z_filtered_normals = filtered_normals[idx]
        return z_filtered_points[:, :3], z_filtered_normals
    
    @staticmethod
    def point_to_hash(point, voxel_size):
        return tuple(np.floor(point / voxel_size).astype(int))