ModelNet40/ModelNet10 数据集预处理

ModelNet数据集是点云处理领域常用的数据集，今天简单分享下预处理ModelNet的方法。
数据集链接：ModelNet40

解压后，数据集内部结构大概就是10&40个分类，内部分为train和test：


打开文件夹能看到，里面的数据格式为off文件。这里我们要转换为txt格式进行下一步处理。

处理代码如下：

import os

import numpy
from tqdm import tqdm
import open3d as o3d
# 文件夹路径
source_path = ''
target_path = ''

# 获取off文件中的xyz数据
def read_off_file(off_file):
    mesh = o3d.io.read_triangle_mesh(off_file)
    # 将OFF文件转换为点云
    point_cloud = numpy.asarray(mesh.vertices)
    return point_cloud


def make_directory(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return
def get_file_name(file_path):
    filename = file_path.split('/')[-1]
    # print(filename) chair_001.off
    (name, suffix) = filename.split('.')
    # print(name, suffix) chair_001 off
    return name

def traversal_directory(path):
    dirs = []
    files = []
    for item in os.scandir(path):
        if item.is_dir():
            dirs.append(item.path)
        elif item.is_file():
            files.append(item.path)
    return dirs, files
def run():
    class_dirs, _ = traversal_directory(source_path)

    for item in class_dirs:
        #bathtub, bed, chair, desk, dresser
        dir_name = item.split('/')[-1]
        # print(dir_name)
        src_train_path = source_path +'/'+dir_name + '/' + 'train'
        src_test_path = source_path +'/'+dir_name + '/' + 'test'
        trg_train_path = target_path +'/'+dir_name + '/' + 'train'
        trg_test_path = target_path +'/'+dir_name + '/' + 'test'
        make_directory(trg_test_path)
        make_directory(trg_train_path)
        _, train_files = traversal_directory(src_train_path)
        _, test_files = traversal_directory(src_test_path)
        for file in train_files:
            name = get_file_name(file)
            pc = read_off_file(file)
            save_path = trg_train_path + '/' + name + '.txt'
            numpy.savetxt(save_path, pc, delimiter=',')
            print(save_path)
        for file in test_files:
        
            name = get_file_name(file)
            pc = read_off_file(file)
            save_path = trg_test_path + '/' + name + '.txt'
            numpy.savetxt(save_path, pc, fmt="%.6f", delimiter=',')
            print(save_path)
if __name__ == "__main__":
    run()

把source_path和target_path换成自己的路径即可。

作者：ミッタン