ModelNet40 中加入自己的数据集

ModelNet40 中加入自己的数据集

ModelNet40 ：http://modelnet.cs.princeton.edu/#
含有40个内别的CAD三维模型，是评价点云深度学习模型进行语意分割、实例分割和分类的标准数据集
代码功能：对自己的点云目标进行标准化融合到ModleNet40中
主要步骤：

1. 中心化：将点云中心移动到坐标原点；
2. 尺度缩放：将所有点的坐标的绝对值限制在1以内；
3. 采样：将点云采样到固定大小点数；
4. 输出保存文件的格式；

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 25 21:53:03 2019

@author: sgl
"""
import os
import sys
import numpy as np
import h5py
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)

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

def loadDataFile(path):
    data = np.loadtxt(path)
    num = data.shape[0]
    point_xyz = data[:,0:3]
    point_normal = data[:,3:6]
    point_rgb = data[:,6:9]
    # label just a example, should be repalced the real.
    # modlenet40 is 0-39, so the label can be 40 and 41
    label = np.ones((num,1), dtype=int)+39
    return point_xyz, label

def change_scale(data):
    #centre 
    xyz_min = np.min(data[:,0:3],axis=0)
    xyz_max = np.max(data[:,0:3],axis=0)
    xyz_move = xyz_min+(xyz_max-xyz_min)/2
    data[:,0:3] = data[:,0:3]-xyz_move
    #scale
    scale = np.max(data[:,0:3])
    data[:,0:3] = data[:,0:3]/scale
    return data

def sample_data(data, num_sample):
    """ data is in N x ...
        we want to keep num_samplexC of them.
        if N > num_sample, we will randomly keep num_sample of them.
        if N < num_sample, we will randomly duplicate samples.
    """
    N = data.shape[0]
    if (N == num_sample):
        return data, range(N)
    elif (N > num_sample):
        sample = np.random.choice(N, num_sample)
        return data[sample, ...], sample
    else:
        sample = np.random.choice(N, num_sample-N)
        dup_data = data[sample, ...]
        return np.concatenate([data, dup_data], 0), list(range(N))+list(sample)
    
if __name__ == "__main__":
    DATA_FILES =getDataFiles(os.path.join(BASE_DIR, 'file_path.txt'))
    num_sample = 4096
    DATA_ALL = []
    for fn in range(len(DATA_FILES)):
        current_data, current_label = loadDataFile(DATA_FILES[fn])
        change_data = change_scale(current_data)
        data_sample,index = sample_data(change_data, num_sample)
        data_label = np.hstack((data_sample,current_label[index]))
        DATA_ALL.append(data_label)
        
    output = np.vstack(DATA_ALL)
    output = output.reshape(-1,num_sample,4)
    
    # train and test number, save data
    if not os.path.exists('plant_train.h5'):
        with h5py.File('plant_train.h5') as f:
            f['data'] = output[0:7,0:3]
            f['labels'] = output[0:8,4]
            
    if not os.path.exists('plant_test.h5'):
        with h5py.File('plant_test.h5') as f:
            f['data'] = output[7:9,0:3]
            f['labels'] = output[7:9,4]

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 25 16:42:25 2019

@author: sgl
"""
import os
import sys
import numpy as np
import h5py
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)

def sample_data(data, num_sample):
   """ data is in N x ...
       we want to keep num_samplexC of them.
       if N > num_sample, we will randomly keep num_sample of them.
       if N < num_sample, we will randomly duplicate samples.
   """
   N = data.shape[0]
   if (N == num_sample):
#        return data, range(N)
       return data
   elif (N > num_sample):
       sample = np.random.choice(N, num_sample)
#        return data[sample, ...], sample
       return data[sample, ...]
   else:
       sample = np.random.choice(N, num_sample-N)
       dup_data = data[sample, ...]
#        return np.concatenate([data, dup_data], 0), list(range(N))+list(sample)
       return np.concatenate([data, dup_data], 0)
   
def creat_pcd_rgba(data,path):
   #write pcd file
#    path = os.path.join(BASE_DIR, 'out_put_'+str(1)+ '.pcd')
   if os.path.exists(path):
       os.remove(path)
   Output_Data = open(path, 'a')
   # headers
   Output_Data.write('# .PCD v0.7 - Point Cloud Data file format\nVERSION 0.7\nFIELDS x y z rgba\nSIZE 4 4 4 4\nTYPE F F F U\nCOUNT 1 1 1 1')
   string = '\nWIDTH ' + str(data.shape[0])
   Output_Data.write(string)
   Output_Data.write('\nHEIGHT 1\nVIEWPOINT 0 0 0 1 0 0 0')
   string = '\nPOINTS ' + str(data.shape[0])
   Output_Data.write(string)
   Output_Data.write('\nDATA ascii')
       
   # pack RGB
   for j in range(data.shape[0]):
       string = ('\n' + str(data[j,0]) + ' ' + str(data[j,1]) + ' ' +str(data[j,2]) + ' ' + str(int(data[j,3])))
       Output_Data.write(string)
       
   Output_Data.close()
   
def creat_pcd_rgb(data,path):
   if os.path.exists(path):
       os.remove(path)
   Output_Data = open(path, 'a')
   # headers
   Output_Data.write('# .PCD v0.7 - Point Cloud Data file format\nVERSION 0.7\nFIELDS x y z rgba\nSIZE 4 4 4 4\nTYPE F F F U\nCOUNT 1 1 1 1')
   string = '\nWIDTH ' + str(data.shape[0])
   Output_Data.write(string)
   Output_Data.write('\nHEIGHT 1\nVIEWPOINT 0 0 0 1 0 0 0')
   string = '\nPOINTS ' + str(data.shape[0])
   Output_Data.write(string)
   Output_Data.write('\nDATA ascii')
       
   # pack RGB
   for j in range(data.shape[0]):
       R=data[j,3]
       G=data[j,4]
       B=data[j,5]
       value = (int(R) << 16 | int(G) << 8 | int(B))
       string = ('\n' + str(data[j,0]) + ' ' + str(data[j,1]) + ' ' +str(data[j,2]) + ' ' + str(value))
       Output_Data.write(string)
       
   Output_Data.close()
   
def creat_txt(data,path):
   np.savetxt(path,data)

data = np.loadtxt("/media/sgl/数据(混合硬盘)/pcl_ubuntu/pcl_test/plant.txt")
#centre 
xyz_min = np.min(data[:,0:3],axis=0)
xyz_max = np.max(data[:,0:3],axis=0)
xyz_move = xyz_min+(xyz_max-xyz_min)/2
data[:,0:3] = data[:,0:3]-xyz_move
#scale
scale = np.max(data[:,0:3])
data[:,0:3] = data[:,0:3]/scale
#sample
data = sample_data(data,4096)
creat_pcd_rgba(data,'/media/sgl/数据(混合硬盘)/PointNet/basisnet/out_put_1.pcd')
creat_txt(data,'/media/sgl/数据(混合硬盘)/PointNet/basisnet/out_put_1.txt')