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数据defread_off_file(off_file):
    mesh = o3d.io.read_triangle_mesh(off_file)# 将OFF文件转换为点云
    point_cloud = numpy.asarray(mesh.vertices)return point_cloud
defmake_directory(path):ifnot os.path.exists(path):
        os.makedirs(path)returndefget_file_name(file_path):
    filename = file_path.split('/')[-1]# print(filename) chair_001.off(name, suffix)= filename.split('.')# print(name, suffix) chair_001 offreturn name
deftraversal_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
defrun():
    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)forfilein 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)forfilein 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换成自己的路径即可。