开源框架PointNet 代码详解——/pointnet/provider.py

provider.py 这个源文件主要用于数据集相关的操作，如数据集增强，数据集加载等。下面依次介绍。

import os
import sys
import numpy as np
import h5py
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)

导入相关库，添加路径。注意这里用到h5py。

# Download dataset for point cloud classification
DATA_DIR = os.path.join(BASE_DIR, 'data')
if not os.path.exists(DATA_DIR):
    os.mkdir(DATA_DIR)
if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')):
    www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip'
    zipfile = os.path.basename(www)
    os.system('wget %s; unzip %s' % (www, zipfile))
    os.system('mv %s %s' % (zipfile[:-4], DATA_DIR))
    os.system('rm %s' % (zipfile))

下载点云分类（3D Object Classification）相关的数据集modelnet40_ply_hdf5_2048。检测到路径下如果没有就自动下载。如果只是做点云语义分割（Semantic Segmentation in Scenes）则不需要下载这一数据集，可以注释掉这一段。

def shuffle_data(data, labels):
    """ Shuffle data and labels.
        Input:
          data: B,N,... numpy array
          label: B,... numpy array
        Return:
          shuffled data, label and shuffle indices
    """
    idx = np.arange(len(labels))
    np.random.shuffle(idx)
    return data[idx, ...], labels[idx], idx

shuffle_data(data, labels) 函数用来在B这个维度上随机打乱数据。注释中输入维度B：batch_size，N：num_points

def rotate_point_cloud(batch_data):
    """ Randomly rotate the point clouds to augument the dataset
        rotation is per shape based along up direction
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        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]])
        shape_pc = batch_data[k, ...]
        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
    return rotated_data

rotate_point_cloud(batch_data) 函数用来绕竖直轴随机旋转点云，作为数据集增强的一部分。因为实际激光雷达采集数据的时候同一个物体由于角度的不同采集到的点云坐标也不同，不希望这导致不同的分类。通过循环旋转batch中的每一个样本（特征向量）。输入点云只有3个维度（XYZ），与旋转变换矩阵相乘返回旋转后的点云。

def rotate_point_cloud_by_angle(batch_data, rotation_angle):

这一函数与上一个类似，只是可以指定旋转角度。

def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
    """ Randomly jitter points. jittering is per point.
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, jittered batch of point clouds
    """
    B, N, C = batch_data.shape
    assert(clip > 0)
    # randn(): standard norm distribution
    # clip(): limit the data within min and max
    jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip)
    jittered_data += batch_data
    return jittered_data

jitter_point_cloud(batch_data, sigma=0.01, clip=0.05) 在原始点云数据集上通过标准正太分布（np.random.randn()）添加噪声，作为数据集增强的一种。噪声数据用np.clip()函数限幅。默认是在点的每一个维度（C：Channel）上添加噪声，也可以按需要修改代码。

def getDataFiles(list_filename):
    return [line.rstrip() for line in open(list_filename)]

def load_h5(h5_filename):
    f = h5py.File(h5_filename)
    data = f['data'][:]
    label = f['label'][:]
    return (data, label)

def loadDataFile(filename):
    return load_h5(filename)

def load_h5_data_label_seg(h5_filename):
    f = h5py.File(h5_filename)
    data = f['data'][:]
    label = f['label'][:]
    seg = f['pid'][:]
    return (data, label, seg)


def loadDataFile_with_seg(filename):
    return load_h5_data_label_seg(filename)

这些函数用来加载h5py格式的数据集