【点云 数据 预处理】制作残缺的点云（一）

参考来源

PF-Net DeCo

思路

比如：全部点[2048,3] = 残缺点[2000,3] +仍掉点[48,3]

方法一：（本文）

残缺点中，没有补其他点，仍然是[2000,3]的大小。（DeCo）

方法二：（下一篇）

残缺点中，填充了[0,0,0]，残缺点云是[2048,3]。（PF-Net）
备注：填充的[0,0,0]点是在原来缺失的位置，
不是机械的在[2000,3]的尾部添加点

文件说明：

以下是两个文件，
第一个是main.py主要执行文件
第二个是shape_utils.py，封装了具体的裁剪的功能

文件1：

main.py 主要执行文件

import torch
import numpy as np
from shape_utils import random_occlude_pointcloud as crop_shape

#  去取点云的txt数据，接收为numpy格式
def readtxt_np(path):
    data_np = np.loadtxt(path)
    return data_np

def np2tensor(data):
    data_tensor = torch.from_numpy(data).float()  #torch.tensor(data)  # torch.float64,即double类型。
    return  data_tensor

#将2维np 转成3维tensor
def np_tensordim3(data_tensor):
    data_tsdim3 = torch.unsqueeze(data_tensor, dim=0)  # 指定位置增加维度
    print(type(data_tsdim3), data_tsdim3.dtype, sep = ' , ')
    return data_tsdim3

if __name__ == '__main__':
    crop_point_num = 500      # 需要丢失的点数
    num_holes = 1             # 固定
    partials,fine_gts= [], []

    # 方式1：读取txt点云数据
    path = "./02691156_2133.txt"
    data_np = readtxt_np(path)
    data_tensor = np2tensor(data_np)
    points = np_tensordim3(data_tensor)

    # # 方式2： 自定义3为维数据
    # points = torch.tensor([[[1, 1, 1], [2, 2, 2], [3, 3, 3],
    #                         [4, 4, 4], [5, 5, 5], [6, 6, 6]],
    #                        [[7, 7, 7], [8, 8, 8], [9, 9, 9],
    #                         [10, 10, 10], [11, 11, 11], [12, 12, 12]]], dtype=torch.float32)
    B, N, dim = points.size()
    print(points.shape)

    # crop centroids： 5 viewpoints to crop around - same crop procedure of PFNet - main paper
    centroids = np.asarray([[0, 0, 1], [1, 0, 0], [1, 0, 1], [-1, 0, 0], [-1, 1, 0]])
    N_partial_points = N - (crop_point_num * num_holes)
    for m in range(B):
        partial, fine_gt = crop_shape(points[m], centroids=centroids,
            scales=[crop_point_num, (crop_point_num )],n_c=num_holes)

        if partial.size(0) > N_partial_points:
            assert num_holes > 1
            # sampling without replacement
            choice = torch.randperm(partial.size(0))[:N_partial_points]
            partial = partial[choice]

        partials.append(partial)
        fine_gts.append(fine_gt)
    partials = torch.stack(partials)  # [B, 3, N-512]
    fine_gts = torch.stack(fine_gts)  # [B, 512, 3]

    print(1111111111111111111111111111111)
    print('残缺：  ', partials.shape)
    print('丢失部分：',fine_gts.shape)

文件2：

shape_utils.py，封装了具体的裁剪的功能

import numpy as np
import open3d as o3
import random
import warnings

def rotate_pointcloud_y(xyz):
    rotation_angle = np.random.uniform() * 2 * np.pi
    cosval = np.cos(rotation_angle)
    sinval = np.sin(rotation_angle)
    rotation_matrix = np.array([[cosval, 0, sinval],
                                [0, 1, 0],
                                [-sinval, 0, cosval]])
    return np.dot(xyz, rotation_matrix)

def rotate_pointcloud_x(xyz):
    rotation_angle = np.random.uniform() * 2 * np.pi
    cosval = np.cos(rotation_angle)
    sinval = np.sin(rotation_angle)
    rotation_matrix = np.array([[1, 0, 0],
                                [0, cosval, -sinval],
                                [0, sinval, cosval]])
    return np.dot(xyz, rotation_matrix)

def jitter_pointcloud(xyz, sigma=0.01, clip=0.05):
    N, C = xyz.shape
    assert (clip > 0)
    jittered_data = np.clip(sigma * np.random.randn(N, C), -1 * clip, clip)
    jittered_data += xyz
    return jittered_data

def center_normalize(pc):
    centroid = np.mean(pc, axis=0)
    pc[:, 0] -= centroid[0]
    pc[:, 1] -= centroid[1]
    pc[:, 2] -= centroid[2]
    d = max(np.sum(np.abs(pc) ** 2, axis=-1) ** (1. / 2))  # furthest_distance
    pc /= d
    return pc

def farthest_point_sample(point, npoint):
    """
    Input:
        xyz: pointcloud data, [N, D]
        npoint: number of samples
    Return:
        fps-point: [npoint, D]
        centroids: [npoint] (indices)
    """
    N, D = point.shape
    xyz = point[:, :3]
    centroids = np.zeros((npoint,))
    distance = np.ones((N,)) * 1e10
    farthest = np.random.randint(0, N)
    for i in range(npoint):
        centroids[i] = farthest
        centroid = xyz[farthest, :]
        dist = np.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = np.argmax(distance, -1)
    centroids = centroids.astype(np.int32)
    point = point[centroids]
    return point, centroids

def random_occlude_pointcloud_v2(xyz, centroids=None, n_drop=200, n_c=1):
    """
    Drops 'n_drop' nearest neighboors around each of 'n_c' centroids
    Parameters
    :param xyz: input pointcloud
    :param centroids: if not None list of keypoints to drop around
    :param n_drop: number of nn to drop around each centroid
    :param n_c: number of centroids selected
    Returns
    :return: partial/cropped point cloud, missing part for completion GT
    """
    in_points = np.shape(xyz)[0]

    if n_drop <= 0:
        # do not crop shape - baseline with no crop augmentation
        return xyz, None

    if centroids is None:
        warnings.warn("None centroids ==> FPS-10")
        _, fps_idx = farthest_point_sample(xyz.numpy(), 10)
        centroids = xyz[fps_idx]

    pcd = o3.geometry.PointCloud()
    pcd.points = o3.utility.Vector3dVector(xyz)
    pcd_tree = o3.geometry.KDTreeFlann(pcd)

    assert len(centroids) >= n_c
    idxs = list(range(len(centroids)))
    random.shuffle(idxs)
    idxs = idxs[:n_c]  # taking only 'n_c' random centroids
    selected_centroids = centroids[idxs]

    dropped = []
    for curr in selected_centroids:
        k, idx, _ = pcd_tree.search_knn_vector_3d(curr, n_drop)  # pcd.points[idx], n_drop)
        dropped.extend(idx)

    accepted_idxs = np.asarray(list(set(list(range(in_points))) - set(dropped)))
    dropped = np.asarray(dropped)
    occluded_pointset = xyz[accepted_idxs]
    missing_part_gt = xyz[dropped]
    return occluded_pointset, missing_part_gt


def random_occlude_pointcloud_v3(xyz, centroids=None, scales=None, n_c=1):
    """

    @param xyz: pointcloud of size N,3
    @param centroids: list of centroids to crop around
    @param scales: [missing part size, missing part size + frame size]
    @param n_c: number of holes
    @return: partial point cloud to complete, missing part gt, missing part + frame gt
    """
    assert scales is not None
    assert len(scales) == 2  # missing + frame
    assert n_c >= 1
    assert len(centroids) >= n_c
    num_points = np.shape(xyz)[0]  # number of input points

    if centroids is None:
        warnings.warn("None centroids")
        _, fps_idx = farthest_point_sample(xyz.numpy(), 10)
        centroids = xyz[fps_idx]

    pcd = o3.geometry.PointCloud()
    pcd.points = o3.utility.Vector3dVector(xyz)
    pcd_tree = o3.geometry.KDTreeFlann(pcd)

    idxs = list(range(len(centroids)))
    random.shuffle(idxs)
    idxs = idxs[:n_c]  # taking only 'n_c' random centroids
    chosen_centroids = centroids[idxs]

    # Scale 0: Missing Part
    dropped = []
    for curr in chosen_centroids:
        k, idx, _ = pcd_tree.search_knn_vector_3d(curr, scales[0])
        dropped.extend(idx)
    idx_accept = np.asarray(list(set(list(range(num_points))) - set(dropped)))
    partial = xyz[idx_accept]  # overall input shape - missing part
    missing_gt = xyz[np.asarray(dropped)]  # missing part GT


    return partial, missing_gt  #, partial_1, missing_gt_1

random_occlude_pointcloud = random_occlude_pointcloud_v3