def s3dis(origin_file_name, label_file_name):
"""
Visualization for S3DIS
"""
origin_data = read_ply(origin_file_name)
label_data = read_ply(label_file_name)
pred = label_data["pred"]
label = label_data["label"]
x = origin_data["x"]
y = origin_data["y"]
z = origin_data["z"]
r = origin_data["red"]
g = origin_data["green"]
b = origin_data["blue"]
label_rgb = np.array([s3dis_color[i] for i in label])
pred_rgb = np.array([s3dis_color[i] for i in pred])
xyzrgb = np.vstack((x, y, z, r, g, b)).T
draw_pc(xyzrgb)
pred_xyzrgb = np.vstack(([x, y, z], pred_rgb.T)).T
draw_pc(pred_xyzrgb)
label_xyzrgb = np.vstack(([x, y, z], label_rgb.T)).T
draw_pc(label_xyzrgb)
def load_sub_sampled_clouds(self, sub_grid_size):
"""load sub_sampled physical files and fill all the containers, input_{trees, colors, labels, names} and val_{proj,labels}-yc
Args:
sub_grid_size ([type]): sub-sampling grid size, e.g., 0.040
"""
tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size))
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.split('/')[-1][:-4]
if self.val_split in cloud_name:
cloud_split = 'validation'
else:
cloud_split = 'training'
# Name of the input files
kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(
cloud_name)) # e.g., Area_1_conferenceRoom_1_KDTree.pkl
sub_ply_file = join(tree_path, '{:s}.ply'.format(
cloud_name)) # e.g., Area_1_conferenceRoom_1.ply
data = read_ply(
sub_ply_file) # ply format: x,y,z,red,gree,blue,class
sub_colors = np.vstack(
(data['red'], data['green'],
data['blue'])).T # (N',3), note the transpose symbol
sub_labels = data['class']
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
# input_xx is a dict contain training or validation info for all sub_pc, each of them contain a list
self.input_trees[cloud_split] += [search_tree]
self.input_colors[cloud_split] += [sub_colors]
self.input_labels[cloud_split] += [sub_labels]
self.input_names[cloud_split] += [cloud_name]
size = sub_colors.shape[0] * 4 * 7
print('{:s} {:.1f} MB loaded in {:.1f}s'.format(
kd_tree_file.split('/')[-1], size * 1e-6,
time.time() - t0))
print('\nPreparing reprojected indices for testing')
# Get validation and test reprojected indices and labels (this is useful for validating on all raw points)
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.split('/')[-1][:-4]
# Validation projection and labels
if self.val_split in cloud_name:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
print('{:s} done in {:.1f}s'.format(cloud_name,
time.time() - t0))
def load_sub_sampled_clouds(self, sub_grid_size):
for cloud in os.listdir(
os.path.join(self.original, 'test', self.test_name)):
print(cloud)
cloud_name = cloud[:-4]
# Name of the input files
kd_tree_file = join(self.sub_folder,
'{:s}_KDTree.pkl'.format(cloud_name))
sub_ply_file = join(self.sub_folder, '{:s}.ply'.format(cloud_name))
sub_data = read_ply(sub_ply_file)
sub_colors = np.vstack(
(sub_data['red'], sub_data['green'], sub_data['blue'])).T
sub_labels = sub_data['class']
full_ply_file = join(self.original, 'test', self.test_name,
'{:s}.ply'.format(cloud_name))
full_data = read_ply(full_ply_file)
full_xyz = np.vstack(
(full_data['x'], full_data['y'], full_data['z'])).T
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
self.input_trees['test'] += [search_tree]
self.input_colors['test'] += [sub_colors]
self.input_labels['test'] += [sub_labels]
self.input_names['test'] += [cloud_name]
self.input_full_xyz['test'] += [full_xyz]
# Test projection indices for testing and labels
proj_file = join(self.sub_folder,
'{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
def load_sub_sampled_clouds(self, sub_grid_size):
tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size))
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.split('/')[-1][:-4]
if self.val_split in cloud_name:
cloud_split = 'validation'
else:
cloud_split = 'training'
# Name of the input files
kd_tree_file = join(tree_path,
'{:s}_KDTree.pkl'.format(cloud_name))
sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name))
data = read_ply(sub_ply_file)
sub_colors = np.vstack(
(data['red'], data['green'], data['blue'])).T
sub_labels = data['class']
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
self.input_trees[cloud_split] += [search_tree]
self.input_colors[cloud_split] += [sub_colors]
self.input_labels[cloud_split] += [sub_labels]
self.input_names[cloud_split] += [cloud_name]
size = sub_colors.shape[0] * 4 * 7
print('{:s} {:.1f} MB loaded in {:.1f}s'.format(
kd_tree_file.split('/')[-1], size * 1e-6,
time.time() - t0))
print('\nPreparing reprojected indices for testing')
# Get validation and test reprojected indices
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.split('/')[-1][:-4]
# Validation projection and labels
if self.val_split in cloud_name:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
print('{:s} done in {:.1f}s'.format(cloud_name,
time.time() - t0))
def read_ply_data(path, with_rgb=True, with_label=True):
data = read_ply(path)
xyz = np.vstack((data['x'], data['y'], data['z'])).T
if with_rgb and with_label:
rgb = np.vstack((data['red'], data['green'], data['blue'])).T
labels = data['class']
return xyz.astype(np.float32), rgb.astype(np.uint8), labels.astype(
np.uint8)
elif with_rgb and not with_label:
rgb = np.vstack((data['red'], data['green'], data['blue'])).T
return xyz.astype(np.float32), rgb.astype(np.uint8)
elif not with_rgb and with_label:
labels = data['class']
return xyz.astype(np.float32), labels.astype(np.uint8)
elif not with_rgb and not with_label:
return xyz.astype(np.float32)
def load_sub_sampled_clouds(self, sub_grid_size):
tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size))
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.replace(
'/', '\\').split('\\')[-1][:-4] # 获取点云名（去除路径去除后缀）
if self.val_split in cloud_name: # 选中的为验证组，其余为训练组
cloud_split = 'val'
else:
cloud_split = 'train'
if cloud_split != self.split: # 跳过非选中组
continue
# Name of the input files
kd_tree_file = join(tree_path,
'{:s}_KDTree.pkl'.format(cloud_name)) # 加载KD树
sub_ply_file = join(tree_path,
'{:s}.ply'.format(cloud_name)) # 加载下采样后的点云数据
data = read_ply(sub_ply_file)
sub_colors = np.vstack(
(data['red'], data['green'], data['blue'])).T # N*3
sub_labels = data['class'] # N*1
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
# self.input_trees += [search_tree.data.astype(np.float32)]
# self.input_colors += [sub_colors.astype(np.float32)]
self.input_feature += [
np.concatenate([
np.array(search_tree.data, copy=False, dtype=np.float32),
sub_colors.astype(np.float32)
],
axis=-1)
]
self.input_labels += [sub_labels.astype(np.int64)]
self.input_names += [cloud_name]
size = sub_colors.shape[0] * 4 * 7
# print(
# '{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.replace('/', '\\').split('\\')[-1], size * 1e-6,
# time.time() - t0))
print('Data Loaded Done!\n') # 为测试准备重投影的指标
def load_sub_sampled_clouds(self, sub_grid_size):
for split in ('train', 'val'):
for cloud in os.listdir(os.path.join(self.original, split)):
cloud_name = cloud[:-4]
cloud_split = split
# Name of the input files
kd_tree_file = join(self.sub_folder,
'{:s}_KDTree.pkl'.format(cloud_name))
sub_ply_file = join(self.sub_folder,
'{:s}.ply'.format(cloud_name))
data = read_ply(sub_ply_file)
sub_colors = np.vstack(
(data['red'], data['green'], data['blue'])).T
sub_labels = data['class']
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
self.input_trees[cloud_split] += [search_tree]
self.input_colors[cloud_split] += [sub_colors]
self.input_labels[cloud_split] += [sub_labels]
self.input_names[cloud_split] += [cloud_name]
size = sub_colors.shape[0] * 4 * 7
print('\nPreparing reprojected indices for testing')
# Get validation and test reprojected indices
for cloud in os.listdir(os.path.join(self.original, "val")):
cloud_name = cloud[:-4]
# Validation projection and labels
proj_file = join(self.sub_folder,
'{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
def load_sub_sampled_clouds(self, sub_grid_size):
for cloud in natsorted(os.listdir(self.path)):
cloud_name = cloud[:-4]
full_ply_file = join(self.path, '{:s}.ply'.format(cloud_name))
full_data = read_ply(full_ply_file)
full_xyz = np.vstack(
(full_data['x'], full_data['y'], full_data['z'])).T
full_colors = np.vstack(
(full_data['red'], full_data['green'], full_data['blue'])).T
full_labels = full_data['class']
xyz_min = np.amin(
full_xyz, axis=0
)[0:
3] # TODO probar sin esto, se ha de haber entrenado sin (data prepare)
full_xyz -= xyz_min
sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(
full_xyz, full_colors, full_labels, sub_grid_size)
sub_colors = sub_colors / 255.0
search_tree = KDTree(sub_xyz)
proj_idx = np.squeeze(
search_tree.query(full_xyz, return_distance=False))
proj_idx = proj_idx.astype(np.int32)
self.input_trees['validation'] += [search_tree]
self.input_colors['validation'] += [sub_colors]
self.input_labels['validation'] += [sub_labels]
self.input_names['validation'] += [cloud_name]
self.input_full_xyz['validation'] += [full_xyz]
self.val_proj += [proj_idx]
self.val_labels += [full_labels]
def semantic3d(origin_file_name, label_file_name):
"""
Visualization for Semantic3D
"""
origin_data = read_ply(origin_file_name)
with open(label_file_name, 'r') as f:
pred = f.readlines()
x = origin_data["x"]
y = origin_data["y"]
z = origin_data["z"]
r = origin_data["red"]
g = origin_data["green"]
b = origin_data["blue"]
pred_rgb = np.array([semantic3d_color[int(i) - 1] for i in pred])
xyzrgb = np.vstack((x, y, z, r, g, b)).T
draw_pc(xyzrgb)
pred_xyzrgb = np.vstack(([x, y, z], pred_rgb.T)).T
draw_pc(pred_xyzrgb)
from helper_ply import read_ply
if __name__ == '__main__':
base_dir = '/data/S3DIS/results'
original_data_dir = '/data/S3DIS/original_ply'
data_path = glob.glob(os.path.join(base_dir, '*.ply'))
data_path = np.sort(data_path)
test_total_correct = 0
test_total_seen = 0
gt_classes = [0 for _ in range(13)]
positive_classes = [0 for _ in range(13)]
true_positive_classes = [0 for _ in range(13)]
for file_name in data_path:
pred_data = read_ply(file_name)
pred = pred_data['pred']
original_data = read_ply(os.path.join(original_data_dir, file_name.split('/')[-1][:-4] + '.ply'))
labels = original_data['class']
points = np.vstack((original_data['x'], original_data['y'], original_data['z'])).T
correct = np.sum(pred == labels)
print(str(file_name.split('/')[-1][:-4]) + '_acc:' + str(correct / float(len(labels))))
test_total_correct += correct
test_total_seen += len(labels)
for j in range(len(labels)):
gt_l = int(labels[j])
pred_l = int(pred[j])
gt_classes[gt_l] += 1
positive_classes[pred_l] += 1
def load_sub_sampled_clouds(self, sub_grid_size):
"""
Load point clouds
:INPUT:
sub_grid_size: Definition of the grid in which is going to be splitted the point cloud
NOTE: Take care with the scale [m,cm,dm] of your point clouds
:RETURN:
-
"""
tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size))
files = np.hstack((self.train_files, self.val_files, self.test_files))
for i, file_path in enumerate(files):
cloud_name = file_path.split('/')[-1][:-4]
if(self.verbose):
print('Load_pc_' + str(i) + ': ' + cloud_name)
if file_path in self.val_files:
cloud_split = 'validation'
elif file_path in self.train_files:
cloud_split = 'train'
else:
cloud_split = 'test'
# Name of the input files
kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name))
sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name))
# read ply with data -- from data preparation script
data = read_ply(sub_ply_file)
sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T
if cloud_split == 'test':
sub_labels = None
else:
sub_labels = data['class']
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
self.input_trees[cloud_split] += [search_tree]
self.input_colors[cloud_split] += [sub_colors]
if cloud_split in ['train', 'validation']:
self.input_labels[cloud_split] += [sub_labels]
# Get validation and test re_projection indices
if(self.verbose):
print('\nPreparing reprojection indices for validation and test')
for i, file_path in enumerate(files):
# get cloud name and split
cloud_name = file_path.split('/')[-1][:-4]
# Validation projection and labels
if file_path in self.val_files:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
# Test projection
if file_path in self.test_files:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.test_proj += [proj_idx]
self.test_labels += [labels]
if(self.verbose):
print('finished')
return
def load_evaluation_points(file_path):
data = read_ply(file_path)
return np.vstack((data['x'], data['y'], data['z'])).T
def load_sub_sampled_clouds(self, sub_grid_size):
tree_path = join(self.testpath, 'input_{:.3f}'.format(sub_grid_size))
files = np.hstack((self.train_files, self.val_files, self.test_files))
for i, file_path in enumerate(files):
cloud_name = file_path.split('/')[-1][:-4]
print('Load_pc_' + str(i) + ': ' + cloud_name)
if file_path in self.val_files:
cloud_split = 'validation'
elif file_path in self.train_files:
cloud_split = 'training'
else:
cloud_split = 'test'
# Name of the input files
kd_tree_file = join(tree_path,
'{:s}_KDTree.pkl'.format(cloud_name))
sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name))
# read ply with data
data = read_ply(sub_ply_file)
#sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T
if cloud_split == 'test':
sub_labels = None
else:
sub_labels = data['class']
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
self.input_trees[cloud_split] += [search_tree]
#self.input_colors[cloud_split] += [sub_colors]
if cloud_split in ['training', 'validation']:
self.input_labels[cloud_split] += [sub_labels]
# Get validation and test re_projection indices
print('\nPreparing reprojection indices for validation and test')
for i, file_path in enumerate(files):
# get cloud name and split
cloud_name = file_path.split('/')[-1][:-4]
print(file_path)
# Validation projection and labels
if file_path in self.val_files:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
# Test projection
if file_path in self.test_files:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.test_proj += [proj_idx]
self.test_labels += [labels]
print('finished')
return
def load_evaluation_rgb(file_path):
data = read_ply(file_path)
return np.vstack((data['red'], data['green'], data['blue'])).T
def convert_pc2plyandweaklabels(anno_path, save_path, sub_pc_folder,
weak_label_folder, weak_label_ratio,
sub_grid_size, gt_class, gt_class2label):
"""
Convert original dataset files (consiting of rooms) to ply file and weak labels. Physically, each room will generate several files, including raw_pc.ply, sub_pc.ply, sub_pc.pkl (for the kdtree), proj_idx.pkl (for each raw point's nearest neighbor in the sub_pc) and weak labels for raw and sub_pc, respectively )
:param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/
:param save_path: path to save original point clouds (each line is XYZRGBL), e.g., xx.ply
:return: None
"""
# save raw_cloud
if not os.path.exists(save_path):
data_list = []
# aggregate a room's instances into 1 pc
for f in glob.glob(join(anno_path, '*.txt')):
class_name = os.path.basename(f).split('_')[0]
if class_name not in gt_class: # note: in some room there is 'staris' class..
class_name = 'clutter'
pc = pd.read_csv(f, header=None, delim_whitespace=True).values
labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name]
data_list.append(np.concatenate([pc, labels], 1)) # Nx7
# translate the data by xyz_min--yc
pc_label = np.concatenate(data_list, 0) # Nx7 as a np object
xyz_min = np.amin(pc_label, axis=0)[0:3]
pc_label[:, 0:3] -= xyz_min
# manage data types and save in PLY format--yc
xyz = pc_label[:, :3].astype(np.float32)
colors = pc_label[:, 3:6].astype(np.uint8)
labels = pc_label[:, 6].astype(np.uint8)
write_ply(save_path, (xyz, colors, labels),
['x', 'y', 'z', 'red', 'green', 'blue', 'class'])
else:
# if existed then read this ply file to fill the data/xyz/colors/labels
data = read_ply(save_path) # ply format: x,y,z,red,gree,blue,class
xyz = np.vstack((data['x'], data['y'],
data['z'])).T # (N',3), note the transpose symbol
colors = np.vstack(
(data['red'], data['green'],
data['blue'])).T # (N',3), note the transpose symbol
labels = data['class']
pc_label = np.concatenate(
(xyz, colors, np.expand_dims(labels, axis=1)), axis=1) # (N,7)
# save sub_cloud
sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply')
if not os.path.exists(sub_ply_file):
sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(
xyz, colors, labels, sub_grid_size)
sub_colors = sub_colors / 255.0
write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels],
['x', 'y', 'z', 'red', 'green', 'blue', 'class'])
else:
data = read_ply(sub_ply_file) # ply format: x,y,z,red,gree,blue,class
sub_xyz = np.vstack((data['x'], data['y'],
data['z'])).T # (N',3), note the transpose symbol
sub_colors = np.vstack(
(data['red'], data['green'],
data['blue'])).T # (N',3), note the transpose symbol
sub_labels = data['class']
# save KDTree for sub_pc
kd_tree_file = join(sub_pc_folder,
str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl')
if not os.path.exists(kd_tree_file):
search_tree = KDTree(sub_xyz)
with open(kd_tree_file, 'wb') as f:
pickle.dump(search_tree, f)
# save projection indcies for all raw points over the corresponding sub_pc
proj_save = join(sub_pc_folder,
str(save_path.split('/')[-1][:-4]) + '_proj.pkl')
if not os.path.exists(proj_save):
proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False))
proj_idx = proj_idx.astype(np.int32)
with open(proj_save, 'wb') as f:
pickle.dump([proj_idx, labels], f)
"""
USED for weakly semantic segmentation
- save raw pc's weak labels
- save sub pc's weak labels
"""
# save raw pc's weak label mask
weak_label_ply_file = join(
weak_label_folder,
save_path.split('/')[-1][:-4] + '_weak_label.ply')
if not os.path.exists(weak_label_ply_file):
# set weak points by randomly selecting weak_label_ratio*N points and denote them w. a mask
num_cloud_points = pc_label.shape[0]
weak_label_mask = np.zeros((num_cloud_points, 1), dtype=np.uint8)
# BUG FIXED: fixed already; here, should set replace = True, otherwise a bug will be resulted
selected_idx = np.random.choice(num_cloud_points,
int(num_cloud_points *
weak_label_ratio),
replace=False)
weak_label_mask[selected_idx, :] = 1
write_ply(weak_label_ply_file, (weak_label_mask, ), ['weak_mask'])
else:
data = read_ply(
weak_label_ply_file) # ply format: x,y,z,red,gree,blue,class
weak_label_mask = data['weak_mask']
# save sub pc's weak label mask
weak_label_sub_file = join(
weak_label_folder,
save_path.split('/')[-1][:-4] + '_sub_weak_label.ply')
if not os.path.exists(weak_label_sub_file):
# HACK: the grid_sub_sampling is deterministic if the inputs are the same. So sub_xyz and sub_colors are the same when called 2nd time
_, _, weak_label_sub_mask = DP.grid_sub_sampling(
xyz, colors, weak_label_mask, sub_grid_size)
write_ply(weak_label_sub_file, (weak_label_sub_mask, ), ['weak_mask'])
def load_sub_sampled_clouds(self, sub_grid_size):
"""load sub_sampled physical files and fill all the containers, input_{trees, colors, labels, names} and val_{proj,labels} and weak label_masks-yc
Args:
sub_grid_size ([type]): sub-sampling grid size, e.g., 0.040
"""
tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size))
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.split('/')[-1][:-4]
if self.val_split in cloud_name:
cloud_split = 'validation'
else:
cloud_split = 'training'
# Name of the input files
kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(
cloud_name)) # e.g., Area_1_conferenceRoom_1_KDTree.pkl
sub_ply_file = join(tree_path, '{:s}.ply'.format(
cloud_name)) # e.g., Area_1_conferenceRoom_1.ply
data = read_ply(
sub_ply_file) # ply format: x,y,z,red,gree,blue,class
sub_colors = np.vstack(
(data['red'], data['green'],
data['blue'])).T # (N',3), note the transpose symbol
sub_labels = data['class']
# read weak labels for sub_pc
weak_label_folder = join(
self.path, 'weak_label_{}'.format(self.weak_label_ratio))
weak_label_sub_file = join(
weak_label_folder,
file_path.split('/')[-1][:-4] + '_sub_weak_label.ply')
if os.path.exists(weak_label_sub_file):
weak_data = read_ply(weak_label_sub_file
) # ply format: x,y,z,red,gree,blue,class
weak_label_sub_mask = weak_data[
'weak_mask'] # (N',) to align same shape as sub_labels
else:
raise NotImplementedError(
"run the dataset_prepare_s3dis_sqn.py to generate weak labels for raw and sub PC"
)
# Read pkl with search tree
with open(kd_tree_file, 'rb') as f:
search_tree = pickle.load(f)
# input_xx is a dict contain training or validation info for all sub_pc, each of them contain a list
self.input_trees[cloud_split] += [search_tree]
self.input_colors[cloud_split] += [sub_colors]
self.input_labels[cloud_split] += [sub_labels]
# HACK: for validation set, all points should have labels meaning the weak_label_ratio is 1(i.e., all points have labels)
if cloud_split == 'validation':
self.input_weak_labels[cloud_split] += [
np.ones_like(weak_label_sub_mask)
]
else:
self.input_weak_labels[cloud_split] += [weak_label_sub_mask]
size = sub_colors.shape[0] * 4 * 7
print('{:s} {:.1f} MB loaded in {:.1f}s'.format(
kd_tree_file.split('/')[-1], size * 1e-6,
time.time() - t0))
print('\nPreparing reprojected indices for testing')
# Get validation and test reprojected indices and labels (this is useful for validating on all raw points)
for i, file_path in enumerate(self.all_files):
t0 = time.time()
cloud_name = file_path.split('/')[-1][:-4]
# Validation projection and labels
if self.val_split in cloud_name:
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
with open(proj_file, 'rb') as f:
proj_idx, labels = pickle.load(f)
self.val_proj += [proj_idx]
self.val_labels += [labels]
print('{:s} done in {:.1f}s'.format(cloud_name,
time.time() - t0))
