点云处理：基于Paddle2.0实现VoxNet对点云进行分类处理

!unzip data/data50045/modelnet40_normal_resampled.zip!mkdir dataset

import osimport numpy as npimport randomimport paddleimport paddle.nn.functional as Ffrom paddle.nn import Conv3D, MaxPool3D, Linear, BatchNorm, Dropout, LeakyReLU, Softmax, Sequential

category = {    'bathtub': 0,    'bed': 1,    'chair': 2,    'door': 3,    'dresser': 4,    'airplane': 5,    'piano': 6,    'sofa': 7,    'person': 8,    'cup': 9}
_category = {    0: 'bathtub',    1: 'bed',    2: 'chair',    3: 'door',    4: 'dresser',    5: 'airplane',    6: 'piano',    7: 'sofa',    8: 'person',    9: 'cup'}
categoryList = [    'bathtub',    'bed',    'chair',    'door',    'dresser',    'airplane',    'piano',    'sofa',    'person',    'cup']

def transform():
    for i in range(len(categoryList)):
        dirpath = os.path.join('./modelnet40_normal_resampled', categoryList[i])
        dirlist = os.listdir(dirpath)        if not os.path.exists(os.path.join('./dataset', categoryList[i])):
            os.mkdir(os.path.join('./dataset', categoryList[i]))        for datalist in dirlist:
            datapath = os.path.join(dirpath, datalist)
            zdata = []
            xdata = []
            ydata = []
            f = open(datapath, 'r')            for point in f:
                xdata.append(float(point.split(',')[0]))
                ydata.append(float(point.split(',')[1]))
                zdata.append(float(point.split(',')[2]))
            f.close()
            arr = np.zeros((32,) * 3).astype('float32')            for j in range(len(xdata)):
                arr[int(xdata[j] * 15.5 + 15.5)][int(ydata[j] * 15.5 + 15.5)][int(zdata[j] * 15.5 + 15.5)] = 1
            savepath = os.path.join('./dataset', categoryList[i], datalist.split('.')[0]+'.npy')
            np.save(savepath, arr)if __name__ == '__main__':
    transform()

def getDatalist():
    f_train = open('./dataset/train.txt', 'w')
    f_test = open('./dataset/test.txt', 'w')    for i in range(len(categoryList)):
        dict_path = os.path.join('./dataset/', categoryList[i])
        data_dict = os.listdir(dict_path)
        count = 0
        for data_path in data_dict:            if count % 60 != 0:
                f_train.write(os.path.join(dict_path, data_path) + ' ' + categoryList[i]+ '\n')            else:
                f_test.write(os.path.join(dict_path, data_path) + ' ' + categoryList[i]+ '\n')
            count += 1
    f_train.close()
    f_test.close()if __name__ == '__main__':
    getDatalist()

def pointDataLoader(file_path='./dataset/train.txt', mode='train'):
    BATCHSIZE = 64
    MAX_POINT = 1024
    datas = []
    labels = []
    f = open(file_path)    for data_list in f:
        point_data = np.load(data_list.split(' ')[0])
        datas.append(point_data)
        labels.append(category[data_list.split(' ')[1].split('\n')[0]])
    f.close()
    datas = np.array(datas)
    labels = np.array(labels)

    index_list = list(range(len(datas)))    def pointDataGenerator():
        if mode == 'train':
            random.shuffle(index_list)
        datas_list = []
        labels_list = []        for i in index_list:
            data = np.reshape(datas[i], [1, 32, 32, 32]).astype('float32')
            label = np.reshape(labels[i], [1]).astype('int64')
            datas_list.append(data) 
            labels_list.append(label)            if len(datas_list) == BATCHSIZE:                yield np.array(datas_list), np.array(labels_list)
                datas_list = []
                labels_list = []        if len(datas_list) > 0:            yield np.array(datas_list), np.array(labels_list)    return pointDataGenerator

class VoxNet(paddle.nn.Layer):
    def __init__(self, name_scope='VoxNet_', num_classes=10):
        super(VoxNet, self).__init__()
        self.backbone = Sequential(
            Conv3D(1, 32, 5, 2),
            BatchNorm(32),
            LeakyReLU(),
            Conv3D(32, 32, 3, 1),
            MaxPool3D(2, 2, 0)
        )
        self.head = Sequential(
            Linear(32*6*6*6, 128),
            LeakyReLU(),
            Dropout(0.2),
            Linear(128, num_classes),
            Softmax()
        )    def forward(self, inputs):
        x = self.backbone(inputs)
        x = paddle.reshape(x, (-1, 32*6*6*6))
        x = self.head(x)        return x

voxnet = VoxNet()
paddle.summary(voxnet, (64, 1, 32, 32, 32))

def train():
    train_loader = pointDataLoader(file_path='./dataset/train.txt', mode='train')
    model = VoxNet()
    optim = paddle.optimizer.Adam(parameters=model.parameters(), weight_decay=0.001)

    loss_fn = paddle.nn.CrossEntropyLoss()
    epoch_num = 100
    for epoch in range(epoch_num):        for batch_id, data in enumerate(train_loader()):
            inputs = paddle.to_tensor(data[0])
            labels = paddle.to_tensor(data[1])
            predicts = model(inputs)

            loss = F.cross_entropy(predicts, labels)
            acc = paddle.metric.accuracy(predicts, labels)
            loss.backward()
            optim.step()
            optim.clear_grad()            if batch_id % 8 == 0: 
                print("epoch: {}, batch_id: {}, loss is: {}, acc is: {}".format(epoch, batch_id, loss.numpy(), acc.numpy()))        if epoch % 20 == 0:
            paddle.save(model.state_dict(), './model/VoxNet.pdparams')
            paddle.save(optim.state_dict(), './model/VoxNet.pdopt')if __name__ == '__main__':
    train()

def evaluation():
    test_loader = pointDataLoader(file_path='./dataset/test.txt', mode='test')
    model = VoxNet()
    model_state_dict = paddle.load('./model/VoxNet.pdparams')
    model.load_dict(model_state_dict)    for batch_id, data in enumerate(test_loader()):
        inputs = paddle.to_tensor(data[0])
        labels = paddle.to_tensor(data[1])
        predicts = model(inputs)

        loss = F.cross_entropy(predicts, labels)
        acc = paddle.metric.accuracy(predicts, labels)        # 打印信息
        if batch_id % 100 == 0:            print("batch_id: {}, loss is: {}, acc is: {}".format(batch_id, loss.numpy(), acc.numpy()))if __name__ == '__main__':
    evaluation()

def visual(file_path='./dataset/predict.txt'):
    f = open(file_path)
    arr = np.load(f.readline().split(' ')[0])
    f.close()
    zdata = []
    xdata = []
    ydata = []    for i in range(arr.shape[0]):        for j in range(arr.shape[1]):            for k in range(arr.shape[2]):                if arr[i][j][k] == 1:
                    xdata.append(float(i))
                    ydata.append(float(j))
                    zdata.append(float(k))
    xdata = np.array(xdata)
    ydata = np.array(ydata)
    zdata = np.array(zdata)    from mpl_toolkits import mplot3d
    %matplotlib inline    import matplotlib.pyplot as plt

    ax = plt.axes(projection='3d')
    ax.scatter3D(xdata, ydata, zdata, c='r')
    plt.show()

def test():
    test_loader = pointDataLoader(file_path='./dataset/predict.txt', mode='test')
    model = VoxNet()
    model_state_dict = paddle.load('./model/VoxNet.pdparams')
    model.load_dict(model_state_dict)    for batch_id, data in enumerate(test_loader()):
        predictdata = paddle.to_tensor(data[0])
        label = paddle.to_tensor(data[1])
        predict = model(predictdata)        print("predict: {}, label: {}".format(np.argmax(predict.numpy(), 1), np.squeeze(label.numpy())))        print("predict: {}".format(_category[int(np.argmax(predict.numpy(), 1))]))if __name__ == '__main__':
    visual()
    test()