forked from charlesq34/pointnet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_prep_util.py
145 lines (124 loc) · 5.07 KB
/
data_prep_util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import os
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
from plyfile import (PlyData, PlyElement, make2d, PlyParseError, PlyProperty)
import numpy as np
import h5py
SAMPLING_BIN = os.path.join(BASE_DIR, 'third_party/mesh_sampling/build/pcsample')
SAMPLING_POINT_NUM = 2048
SAMPLING_LEAF_SIZE = 0.005
MODELNET40_PATH = '../datasets/modelnet40'
def export_ply(pc, filename):
vertex = np.zeros(pc.shape[0], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')])
for i in range(pc.shape[0]):
vertex[i] = (pc[i][0], pc[i][1], pc[i][2])
ply_out = PlyData([PlyElement.describe(vertex, 'vertex', comments=['vertices'])])
ply_out.write(filename)
# Sample points on the obj shape
def get_sampling_command(obj_filename, ply_filename):
cmd = SAMPLING_BIN ' ' obj_filename
cmd = ' ' ply_filename
cmd = ' -n_samples %d ' % SAMPLING_POINT_NUM
cmd = ' -leaf_size %f ' % SAMPLING_LEAF_SIZE
return cmd
# --------------------------------------------------------------
# Following are the helper functions to load MODELNET40 shapes
# --------------------------------------------------------------
# Read in the list of categories in MODELNET40
def get_category_names():
shape_names_file = os.path.join(MODELNET40_PATH, 'shape_names.txt')
shape_names = [line.rstrip() for line in open(shape_names_file)]
return shape_names
# Return all the filepaths for the shapes in MODELNET40
def get_obj_filenames():
obj_filelist_file = os.path.join(MODELNET40_PATH, 'filelist.txt')
obj_filenames = [os.path.join(MODELNET40_PATH, line.rstrip()) for line in open(obj_filelist_file)]
print('Got %d obj files in modelnet40.' % len(obj_filenames))
return obj_filenames
# Helper function to create the father folder and all subdir folders if not exist
def batch_mkdir(output_folder, subdir_list):
if not os.path.exists(output_folder):
os.mkdir(output_folder)
for subdir in subdir_list:
if not os.path.exists(os.path.join(output_folder, subdir)):
os.mkdir(os.path.join(output_folder, subdir))
# ----------------------------------------------------------------
# Following are the helper functions to load save/load HDF5 files
# ----------------------------------------------------------------
# Write numpy array data and label to h5_filename
def save_h5_data_label_normal(h5_filename, data, label, normal,
data_dtype='float32', label_dtype='uint8', noral_dtype='float32'):
h5_fout = h5py.File(h5_filename)
h5_fout.create_dataset(
'data', data=data,
compression='gzip', compression_opts=4,
dtype=data_dtype)
h5_fout.create_dataset(
'normal', data=normal,
compression='gzip', compression_opts=4,
dtype=normal_dtype)
h5_fout.create_dataset(
'label', data=label,
compression='gzip', compression_opts=1,
dtype=label_dtype)
h5_fout.close()
# Write numpy array data and label to h5_filename
def save_h5(h5_filename, data, label, data_dtype='uint8', label_dtype='uint8'):
h5_fout = h5py.File(h5_filename)
h5_fout.create_dataset(
'data', data=data,
compression='gzip', compression_opts=4,
dtype=data_dtype)
h5_fout.create_dataset(
'label', data=label,
compression='gzip', compression_opts=1,
dtype=label_dtype)
h5_fout.close()
# Read numpy array data and label from h5_filename
def load_h5_data_label_normal(h5_filename):
f = h5py.File(h5_filename)
data = f['data'][:]
label = f['label'][:]
normal = f['normal'][:]
return (data, label, normal)
# Read numpy array data and label from 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)
# Read numpy array data and label from h5_filename
def load_h5(h5_filename):
f = h5py.File(h5_filename)
data = f['data'][:]
label = f['label'][:]
return (data, label)
# ----------------------------------------------------------------
# Following are the helper functions to load save/load PLY files
# ----------------------------------------------------------------
# Load PLY file
def load_ply_data(filename, point_num):
plydata = PlyData.read(filename)
pc = plydata['vertex'].data[:point_num]
pc_array = np.array([[x, y, z] for x,y,z in pc])
return pc_array
# Load PLY file
def load_ply_normal(filename, point_num):
plydata = PlyData.read(filename)
pc = plydata['normal'].data[:point_num]
pc_array = np.array([[x, y, z] for x,y,z in pc])
return pc_array
# Make up rows for Nxk array
# Input Pad is 'edge' or 'constant'
def pad_arr_rows(arr, row, pad='edge'):
assert(len(arr.shape) == 2)
assert(arr.shape[0] <= row)
assert(pad == 'edge' or pad == 'constant')
if arr.shape[0] == row:
return arr
if pad == 'edge':
return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'edge')
if pad == 'constant':
return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'constant', (0, 0))