def __getitem__(self, idx):
"""
Returns a dict with following keys:
point_clouds: (N,3+C)
center_label: (MAX_NUM_OBJ,3) for GT box center XYZ
sem_cls_label: (MAX_NUM_OBJ,) semantic class index
angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1
angle_residual_label: (MAX_NUM_OBJ,)
size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER
size_residual_label: (MAX_NUM_OBJ,3)
box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box
point_votes: (N,3) with votes XYZ
point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB.
scan_idx: int scan index in scan_names list
pcl_color: unused
"""
scan_name = self.scan_names[idx]
mesh_vertices = np.load(
os.path.join(self.data_path, scan_name) + '_vert.npy')
instance_labels = np.load(
os.path.join(self.data_path, scan_name) + '_ins_label.npy')
semantic_labels = np.load(
os.path.join(self.data_path, scan_name) + '_sem_label.npy')
instance_bboxes = np.load(
os.path.join(self.data_path, scan_name) + '_bbox.npy')
if not self.use_color:
point_cloud = mesh_vertices[:, 0:3] # do not use color for now
pcl_color = mesh_vertices[:, 3:6]
else:
point_cloud = mesh_vertices[:, 0:6]
point_cloud[:, 3:] = (point_cloud[:, 3:] - MEAN_COLOR_RGB) / 256.0
if self.use_height:
floor_height = np.percentile(point_cloud[:, 2], 0.99)
height = point_cloud[:, 2] - floor_height
point_cloud = np.concatenate(
[point_cloud, np.expand_dims(height, 1)], 1)
# ------------------------------- LABELS ------------------------------
target_bboxes = np.zeros((MAX_NUM_OBJ, 6))
target_bboxes_mask = np.zeros((MAX_NUM_OBJ))
angle_classes = np.zeros((MAX_NUM_OBJ, ))
angle_residuals = np.zeros((MAX_NUM_OBJ, ))
size_classes = np.zeros((MAX_NUM_OBJ, ))
size_residuals = np.zeros((MAX_NUM_OBJ, 3))
point_cloud, choices = pc_util.random_sampling(point_cloud,
self.num_points,
return_choices=True)
instance_labels = instance_labels[choices]
semantic_labels = semantic_labels[choices]
pcl_color = pcl_color[choices]
target_bboxes_mask[0:instance_bboxes.shape[0]] = 1
target_bboxes[0:instance_bboxes.shape[0], :] = instance_bboxes[:, 0:6]
# ------------------------------- DATA AUGMENTATION ------------------------------
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
point_cloud[:, 0] = -1 * point_cloud[:, 0]
target_bboxes[:, 0] = -1 * target_bboxes[:, 0]
if np.random.random() > 0.5:
# Flipping along the XZ plane
point_cloud[:, 1] = -1 * point_cloud[:, 1]
target_bboxes[:, 1] = -1 * target_bboxes[:, 1]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random() * np.pi /
18) - np.pi / 36 # -5 ~ +5 degree
rot_mat = pc_util.rotz(rot_angle)
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat)
# compute votes *AFTER* augmentation
# generate votes
# Note: since there's no map between bbox instance labels and
# pc instance_labels (it had been filtered
# in the data preparation step) we'll compute the instance bbox
# from the points sharing the same instance label.
point_votes = np.zeros([self.num_points, 3])
point_votes_mask = np.zeros(self.num_points)
for i_instance in np.unique(instance_labels):
# find all points belong to that instance
ind = np.where(instance_labels == i_instance)[0]
# find the semantic label
if semantic_labels[ind[0]] in DC.nyu40ids:
x = point_cloud[ind, :3]
center = 0.5 * (x.min(0) + x.max(0))
point_votes[ind, :] = center - x
point_votes_mask[ind] = 1.0
point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical
class_ind = [
np.where(DC.nyu40ids == x)[0][0] for x in instance_bboxes[:, -1]
]
# NOTE: set size class as semantic class. Consider use size2class.
size_classes[0:instance_bboxes.shape[0]] = class_ind
size_residuals[0:instance_bboxes.shape[0], :] = \
target_bboxes[0:instance_bboxes.shape[0], 3:6] - DC.mean_size_arr[class_ind,:]
# keep the same nums of points for each cloud
mesh_vertices, _ = pc_util.random_sampling(mesh_vertices,
50000,
return_choices=True)
ret_dict = {}
ret_dict['mesh_vertices'] = mesh_vertices.astype(np.float32)
ret_dict['point_clouds'] = point_cloud.astype(np.float32)
ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3]
ret_dict['heading_class_label'] = angle_classes.astype(np.int64)
ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32)
ret_dict['size_class_label'] = size_classes.astype(np.int64)
ret_dict['size_residual_label'] = size_residuals.astype(np.float32)
target_bboxes_semcls = np.zeros((MAX_NUM_OBJ))
target_bboxes_semcls[0:instance_bboxes.shape[0]] = \
[DC.nyu40id2class[x] for x in instance_bboxes[:,-1][0:instance_bboxes.shape[0]]]
ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64)
ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32)
ret_dict['vote_label'] = point_votes.astype(np.float32)
ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['pcl_color'] = pcl_color
return ret_dict
def __getitem__(self, idx):
"""
Returns a dict with following keys:
point_clouds: (N,3+C)
scan_idx: int scan index in scan_names list
"""
scan_name = self.scan_names[idx]
mesh_vertices = np.load(
os.path.join(self.data_path, scan_name) + '_pc.npz')['pc'] # Nx6
if not self.use_color:
raw_point_cloud = mesh_vertices[:, 0:3] # do not use color for now
else:
raw_point_cloud = mesh_vertices[:, 0:6]
raw_point_cloud[:, 3:] = (raw_point_cloud[:, 3:] -
MEAN_COLOR_RGB) / 256.0
if self.use_height:
floor_height = np.percentile(raw_point_cloud[:, 2], 0.99)
height = raw_point_cloud[:, 2] - floor_height
raw_point_cloud = np.concatenate(
[raw_point_cloud, np.expand_dims(height, 1)], 1)
point_cloud, choices = pc_util.random_sampling(raw_point_cloud,
self.num_points,
return_choices=True)
#ema_point_cloud = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=False)
ema_point_cloud = point_cloud.copy() # 2021.2.28
# ------------------------------- DATA AUGMENTATION ------------------------------
flip_x_axis = 0
flip_y_axis = 0
flip_x_axis_ema = 0 # 2021.2.28
flip_y_axis_ema = 0 # 2021.2.28
rot_mat = np.identity(3)
scale_ratio = np.ones((1, 3))
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
flip_x_axis = 1
point_cloud[:, 0] = -1 * point_cloud[:, 0]
if np.random.random() > 0.5: # 2021.2.28
# Flipping along the YZ plane
flip_x_axis_ema = 1
ema_point_cloud[:, 0] = -1 * ema_point_cloud[:, 0]
if np.random.random() > 0.5:
# Flipping along the XZ plane
flip_y_axis = 1
point_cloud[:, 1] = -1 * point_cloud[:, 1]
if np.random.random() > 0.5: # 2021.2.28
# Flipping along the XZ plane
flip_y_axis_ema = 1
ema_point_cloud[:, 1] = -1 * ema_point_cloud[:, 1]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random() * np.pi /
3) - np.pi / 6 # -30 ~ +30 degree
rot_mat = pc_util.rotz(rot_angle)
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
# Augment point cloud scale: 0.85x-1.15x
scale_ratio = np.random.random() * 0.3 + 0.85
scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0)
point_cloud[:, 0:3] *= scale_ratio
if self.use_height:
point_cloud[:, -1] *= scale_ratio[0, 0]
ret_dict = {}
ret_dict['point_clouds'] = point_cloud.astype(np.float32)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['supervised_mask'] = np.array(0).astype(np.int64)
ret_dict['ema_point_clouds'] = ema_point_cloud.astype(np.float32)
ret_dict['flip_x_axis'] = np.array(flip_x_axis).astype(np.int64)
ret_dict['flip_y_axis'] = np.array(flip_y_axis).astype(np.int64)
ret_dict['rot_mat'] = rot_mat.astype(np.float32)
ret_dict['scale'] = np.array(scale_ratio).astype(np.float32)
ret_dict['flip_x_axis_ema'] = np.array(flip_x_axis_ema).astype(
np.int64) # 2021.2.28
ret_dict['flip_y_axis_ema'] = np.array(flip_y_axis_ema).astype(
np.int64) # 2021.2.28
return ret_dict
def __getitem__(self, idx):
"""
Returns a dict with following keys:
point_clouds: (N,3+C)
center_label: (MAX_NUM_OBJ,3) for GT box center XYZ
sem_cls_label: (MAX_NUM_OBJ,) semantic class index
angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1
angle_residual_label: (MAX_NUM_OBJ,)
size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER
size_residual_label: (MAX_NUM_OBJ,3)
box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box
point_votes: (N,3) with votes XYZ
point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB.
scan_idx: int scan index in scan_names list
pcl_color: unused
"""
scan_name = self.scan_names[idx]
mesh_vertices = np.load(os.path.join(self.data_path, scan_name)+'_vert.npy')
meta_vertices = np.load(os.path.join(self.data_path, scan_name)+'_all_noangle_40cls.npy') ### Need to change the name here
instance_labels = meta_vertices[:,-2]
semantic_labels = meta_vertices[:,-1]
if not self.use_color:
point_cloud = mesh_vertices[:,0:3] # do not use color for now
pcl_color = mesh_vertices[:,3:6]
else:
point_cloud = mesh_vertices[:,0:6]
point_cloud[:,3:] = (point_cloud[:,3:]-MEAN_COLOR_RGB)/256.0
pcl_color = (point_cloud[:,3:]-MEAN_COLOR_RGB)/256.0
if self.use_height:
floor_height = np.percentile(point_cloud[:,2],0.99)
height = point_cloud[:,2] - floor_height
point_cloud = np.concatenate([point_cloud, np.expand_dims(height, 1)],1)
# ------------------------------- LABELS ------------------------------
target_bboxes = np.zeros((MAX_NUM_OBJ, 6))
target_bboxes_mask = np.zeros((MAX_NUM_OBJ))
angle_classes = np.zeros((MAX_NUM_OBJ,))
angle_label = np.zeros((MAX_NUM_OBJ,))
angle_residuals = np.zeros((MAX_NUM_OBJ,))
size_classes = np.zeros((MAX_NUM_OBJ,))
size_residuals = np.zeros((MAX_NUM_OBJ, 3))
### For statistics
surface_cue = np.zeros((MAX_NUM_OBJ))
line_cue = np.zeros((MAX_NUM_OBJ,))
before_sample = np.unique(instance_labels)
while True:
orig_point_cloud = np.copy(point_cloud)
temp_point_cloud, choices = pc_util.random_sampling(orig_point_cloud,
self.num_points, return_choices=True)
after_sample = np.unique(instance_labels[choices])
if np.array_equal(before_sample, after_sample):
point_cloud = temp_point_cloud
break
instance_labels = instance_labels[choices]
semantic_labels = semantic_labels[choices]
meta_vertices = meta_vertices[choices]
pcl_color = pcl_color[choices]
# ------------------------------- DATA AUGMENTATION ------------------------------
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
point_cloud[:,0] = -1 * point_cloud[:,0]
# target_bboxes[:,0] = -1 * target_bboxes[:,0]
meta_vertices[:, 0] = -1 * meta_vertices[:, 0]
meta_vertices[:, 6] = -1 * meta_vertices[:, 6]
if np.random.random() > 0.5:
# Flipping along the XZ plane
point_cloud[:,1] = -1 * point_cloud[:,1]
# target_bboxes[:,1] = -1 * target_bboxes[:,1]
meta_vertices[:, 1] = -1 * meta_vertices[:, 1]
meta_vertices[:, 6] = -1 * meta_vertices[:, 6]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random()*np.pi/18) - np.pi/36 # -5 ~ +5 degree
rot_mat = pc_util.rotz(rot_angle).astype(np.float32)
point_cloud[:,0:3] = np.dot(point_cloud[:,0:3], np.transpose(rot_mat))
meta_vertices[:, :6] = rotate_aligned_boxes(meta_vertices[:, :6], rot_mat)
meta_vertices[:, 6] += rot_angle
# ------------------------------- Plane and point ------------------------------
# compute votes *AFTER* augmentation
# generate votes
# Note: since there's no map between bbox instance labels and
# pc instance_labels (it had been filtered
# in the data preparation step) we'll compute the instance bbox
# from the points sharing the same instance label.
point_votes = np.zeros([self.num_points, 3])
point_votes_mask = np.zeros(self.num_points)
point_boundary_mask_z = np.zeros(self.num_points)
point_boundary_mask_xy = np.zeros(self.num_points)
point_boundary_offset_z = np.zeros([self.num_points, 3])
point_boundary_offset_xy = np.zeros([self.num_points, 3])
point_boundary_sem_z = np.zeros([self.num_points, 3+2+1])
point_boundary_sem_xy = np.zeros([self.num_points, 3+1+1])
point_line_mask = np.zeros(self.num_points)
point_line_offset = np.zeros([self.num_points, 3])
point_line_sem = np.zeros([self.num_points, 3+1])
point_sem_label = np.zeros(self.num_points)
selected_instances = []
selected_centers = []
selected_centers_support = []
selected_centers_bsupport = []
obj_meta = []
counter = -1
for i_instance in np.unique(instance_labels):
# find all points belong to that instance
ind = np.where(instance_labels == i_instance)[0]
if semantic_labels[ind[0]] in DC.nyu40ids:
counter += 1
idx_instance = counter
x = point_cloud[ind,:3]
### Meta information here
meta = meta_vertices[ind[0]]
obj_meta.append(meta)
### Get the centroid here
center = meta[:3]
point_votes[ind, :] = center - x
point_votes_mask[ind] = 1.0
point_sem_label[ind] = DC.nyu40id2class_sem[meta[-1]]
### Corners
corners, xmin, ymin, zmin, xmax, ymax, zmax = params2bbox(center, meta[3], meta[4], meta[5], meta[6])
## Get lower four lines
plane_lower_temp = np.array([0,0,1,-corners[6,-1]])
para_points = np.array([corners[1], corners[3], corners[5], corners[7]])
newd = np.sum(para_points * plane_lower_temp[:3], 1)
if check_upright(para_points) and plane_lower_temp[0]+plane_lower_temp[1] < LOWER_THRESH:
plane_lower = np.array([0,0,1,plane_lower_temp[-1]])
plane_upper = np.array([0,0,1,-np.mean(newd)])
else:
import pdb;pdb.set_trace()
print ("error with upright")
if check_z(plane_upper, para_points) == False:
import pdb;pdb.set_trace()
### Get the boundary points here
alldist = np.abs(np.sum(x*plane_lower[:3], 1) + plane_lower[-1])
mind = np.min(alldist)
sel = np.abs(alldist - mind) < DIST_THRESH
## Get lower four lines
line_sel1, line_sel2, line_sel3, line_sel4 = get_linesel(x[sel], xmin, xmax, ymin, ymax)
if np.sum(line_sel1) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel1]] = 1.0
linecenter = np.mean(x[sel][line_sel1], axis=0)
linecenter[1] = (ymin+ymax)/2.0
point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1]
point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel2) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel2]] = 1.0
linecenter = np.mean(x[sel][line_sel2], axis=0)
linecenter[1] = (ymin+ymax)/2.0
point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2]
point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel3) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel3]] = 1.0
linecenter = np.mean(x[sel][line_sel3], axis=0)
linecenter[0] = (xmin+xmax)/2.0
point_line_offset[ind[sel][line_sel3]] = linecenter - x[sel][line_sel3]
point_line_sem[ind[sel][line_sel3]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel4) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel4]] = 1.0
linecenter = np.mean(x[sel][line_sel4], axis=0)
linecenter[0] = (xmin+xmax)/2.0
point_line_offset[ind[sel][line_sel4]] = linecenter - x[sel][line_sel4]
point_line_sem[ind[sel][line_sel4]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
### Set the surface labels here
if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH:
center = np.array([(xmin+xmax)/2.0, (ymin+ymax)/2.0, np.mean(x[sel][:,2])])
sel_global = ind[sel]
point_boundary_mask_z[sel_global] = 1.0
point_boundary_sem_z[sel_global] = np.array([center[0], center[1], center[2], xmax - xmin, ymax - ymin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
point_boundary_offset_z[sel_global] = center - x[sel]
### Get the boundary points here
alldist = np.abs(np.sum(x*plane_upper[:3], 1) + plane_upper[-1])
mind = np.min(alldist)
sel = np.abs(alldist - mind) < DIST_THRESH
## Get upper four lines
line_sel1, line_sel2, line_sel3, line_sel4 = get_linesel(x[sel], xmin, xmax, ymin, ymax)
if np.sum(line_sel1) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel1]] = 1.0
linecenter = np.mean(x[sel][line_sel1], axis=0)
linecenter[1] = (ymin+ymax)/2.0
point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1]
point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel2) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel2]] = 1.0
linecenter = np.mean(x[sel][line_sel2], axis=0)
linecenter[1] = (ymin+ymax)/2.0
point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2]
point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel3) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel3]] = 1.0
linecenter = np.mean(x[sel][line_sel3], axis=0)
linecenter[0] = (xmin+xmax)/2.0
point_line_offset[ind[sel][line_sel3]] = linecenter - x[sel][line_sel3]
point_line_sem[ind[sel][line_sel3]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel4) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel4]] = 1.0
linecenter = np.mean(x[sel][line_sel4], axis=0)
linecenter[0] = (xmin+xmax)/2.0
point_line_offset[ind[sel][line_sel4]] = linecenter - x[sel][line_sel4]
point_line_sem[ind[sel][line_sel4]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH:
center = np.array([(xmin+xmax)/2.0, (ymin+ymax)/2.0, np.mean(x[sel][:,2])])
sel_global = ind[sel]
point_boundary_mask_z[sel_global] = 1.0
point_boundary_sem_z[sel_global] = np.array([center[0], center[1], center[2], xmax - xmin, ymax - ymin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
point_boundary_offset_z[sel_global] = center - x[sel]
## Get left two lines
v1 = corners[3] - corners[2]
v2 = corners[2] - corners[0]
cp = np.cross(v1, v2)
d = -np.dot(cp,corners[0])
a,b,c = cp
plane_left_temp = np.array([a, b, c, d])
para_points = np.array([corners[4], corners[5], corners[6], corners[7]])
### Normalize xy here
plane_left_temp /= np.linalg.norm(plane_left_temp[:3])
newd = np.sum(para_points * plane_left_temp[:3], 1)
if plane_left_temp[2] < LOWER_THRESH:
plane_left = plane_left_temp#np.array([cls,res,tempsign,plane_left_temp[-1]])
plane_right = np.array([plane_left_temp[0], plane_left_temp[1], plane_left_temp[2], -np.mean(newd)])
else:
import pdb;pdb.set_trace()
print ("error with upright")
### Get the boundary points here
alldist = np.abs(np.sum(x*plane_left[:3], 1) + plane_left[-1])
mind = np.min(alldist)
sel = np.abs(alldist - mind) < DIST_THRESH
## Get upper four lines
line_sel1, line_sel2 = get_linesel2(x[sel], ymin, ymax, zmin, zmax, axis=1)
if np.sum(line_sel1) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel1]] = 1.0
linecenter = np.mean(x[sel][line_sel1], axis=0)
linecenter[2] = (zmin+zmax)/2.0
point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1]
point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel2) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel2]] = 1.0
linecenter = np.mean(x[sel][line_sel2], axis=0)
linecenter[2] = (zmin+zmax)/2.0
point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2]
point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH:
center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0])
sel_global = ind[sel]
point_boundary_mask_xy[sel_global] = 1.0
point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
point_boundary_offset_xy[sel_global] = center - x[sel]
### Get the boundary points here
alldist = np.abs(np.sum(x*plane_right[:3], 1) + plane_right[-1])
mind = np.min(alldist)
sel = np.abs(alldist - mind) < DIST_THRESH
line_sel1, line_sel2 = get_linesel2(x[sel], ymin, ymax, zmin, zmax, axis=1)
if np.sum(line_sel1) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel1]] = 1.0
linecenter = np.mean(x[sel][line_sel1], axis=0)
linecenter[2] = (zmin+zmax)/2.0
point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1]
point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(line_sel2) > NUM_POINT_LINE:
point_line_mask[ind[sel][line_sel2]] = 1.0
linecenter = np.mean(x[sel][line_sel2], axis=0)
linecenter[2] = (zmin+zmax)/2.0
point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2]
point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH:
center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0])
sel_global = ind[sel]
point_boundary_mask_xy[sel_global] = 1.0
point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
point_boundary_offset_xy[sel_global] = center - x[sel]
### Get the boundary points here
v1 = corners[0] - corners[4]
v2 = corners[4] - corners[5]
cp = np.cross(v1, v2)
d = -np.dot(cp,corners[5])
a,b,c = cp
plane_front_temp = np.array([a, b, c, d])
para_points = np.array([corners[2], corners[3], corners[6], corners[7]])
plane_front_temp /= np.linalg.norm(plane_front_temp[:3])
newd = np.sum(para_points * plane_front_temp[:3], 1)
if plane_front_temp[2] < LOWER_THRESH:
plane_front = plane_front_temp
plane_back = np.array([plane_front_temp[0], plane_front_temp[1], plane_front_temp[2], -np.mean(newd)])
else:
import pdb;pdb.set_trace()
print ("error with upright")
### Get the boundary points here
alldist = np.abs(np.sum(x*plane_front[:3], 1) + plane_front[-1])
mind = np.min(alldist)
sel = np.abs(alldist - mind) < DIST_THRESH
if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH:
center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0])
sel_global = ind[sel]
point_boundary_mask_xy[sel_global] = 1.0
point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
point_boundary_offset_xy[sel_global] = center - x[sel]
### Get the boundary points here
alldist = np.abs(np.sum(x*plane_back[:3], 1) + plane_back[-1])
mind = np.min(alldist)
sel = np.abs(alldist - mind) < DIST_THRESH
if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH:
center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0])
sel_global = ind[sel]
point_boundary_mask_xy[sel_global] = 1.0
point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]])
point_boundary_offset_xy[sel_global] = center - x[sel]
num_instance = len(obj_meta)
obj_meta = np.array(obj_meta)
obj_meta = obj_meta.reshape(-1, 9)
target_bboxes_mask[0:num_instance] = 1
target_bboxes[0:num_instance,:6] = obj_meta[:,0:6]
class_ind = [np.where(DC.nyu40ids == x)[0][0] for x in obj_meta[:,-1]]
# NOTE: set size class as semantic class. Consider use size2class.
size_classes[0:num_instance] = class_ind
size_residuals[0:num_instance, :] = \
target_bboxes[0:num_instance, 3:6] - DC.mean_size_arr[class_ind,:]
point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical
point_sem_label = np.tile(np.expand_dims(point_sem_label, -1), (1, 3)) # make 3 votes identical
ret_dict = {}
ret_dict['point_clouds'] = point_cloud.astype(np.float32)
ret_dict['center_label'] = target_bboxes.astype(np.float32)[:,0:3]
ret_dict['size_label'] = target_bboxes.astype(np.float32)[:,3:6]
ret_dict['heading_label'] = angle_label.astype(np.float32)
ret_dict['heading_class_label'] = angle_classes.astype(np.int64)
ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32)
ret_dict['size_class_label'] = size_classes.astype(np.int64)
ret_dict['size_residual_label'] = size_residuals.astype(np.float32)
if self.use_height:
ret_dict['floor_height'] = floor_height
target_bboxes_semcls = np.zeros((MAX_NUM_OBJ))
target_bboxes_semcls[0:num_instance] = \
[DC.nyu40id2class[x] for x in obj_meta[:,-1][0:obj_meta.shape[0]]]
ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64)
ret_dict['point_sem_cls_label'] = point_sem_label.astype(np.int64)
ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32)
ret_dict['point_boundary_mask_z'] = point_boundary_mask_z.astype(np.float32)
ret_dict['point_boundary_mask_xy'] = point_boundary_mask_xy.astype(np.float32)
ret_dict['point_boundary_offset_z'] = point_boundary_offset_z.astype(np.float32)
ret_dict['point_boundary_offset_xy'] = point_boundary_offset_xy.astype(np.float32)
ret_dict['point_boundary_sem_z'] = point_boundary_sem_z.astype(np.float32)
ret_dict['point_boundary_sem_xy'] = point_boundary_sem_xy.astype(np.float32)
ret_dict['point_line_mask'] = point_line_mask.astype(np.float32)
ret_dict['point_line_offset'] = point_line_offset.astype(np.float32)
ret_dict['point_line_sem'] = point_line_sem.astype(np.float32)
ret_dict['vote_label'] = point_votes.astype(np.float32)
ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['pcl_color'] = pcl_color
ret_dict['num_instance'] = num_instance
return ret_dict
def __getitem__(self, idx):
"""
Returns a dict with following keys:
point_clouds: (N,3+C)
center_label: (MAX_NUM_OBJ,3) for GT box center XYZ
sem_cls_label: (MAX_NUM_OBJ,) semantic class index
angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1
angle_residual_label: (MAX_NUM_OBJ,)
size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER
size_residual_label: (MAX_NUM_OBJ,3)
box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box
point_votes: (N,3) with votes XYZ
point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB.
scan_idx: int scan index in scan_names list
"""
scan_name = self.scan_names[idx]
mesh_vertices = np.load(
os.path.join(self.data_path, scan_name) + '_vert.npy')
instance_labels = np.load(
os.path.join(self.data_path, scan_name) + '_ins_label.npy')
semantic_labels = np.load(
os.path.join(self.data_path, scan_name) + '_sem_label.npy')
instance_bboxes = np.load(
os.path.join(self.data_path, scan_name) + '_bbox.npy')
if not self.use_color:
raw_point_cloud = mesh_vertices[:, 0:3] # do not use color for now
else:
raw_point_cloud = mesh_vertices[:, 0:6]
raw_point_cloud[:, 3:] = (raw_point_cloud[:, 3:] -
MEAN_COLOR_RGB) / 256.0
if self.use_height:
floor_height = np.percentile(raw_point_cloud[:, 2], 0.99)
height = raw_point_cloud[:, 2] - floor_height
raw_point_cloud = np.concatenate(
[raw_point_cloud, np.expand_dims(height, 1)], 1)
# ------------------------------- LABELS ------------------------------
target_bboxes = np.zeros((MAX_NUM_OBJ, 6))
target_bboxes_mask = np.zeros((MAX_NUM_OBJ))
angle_classes = np.zeros((MAX_NUM_OBJ, ))
angle_residuals = np.zeros((MAX_NUM_OBJ, ))
size_classes = np.zeros((MAX_NUM_OBJ, ))
size_residuals = np.zeros((MAX_NUM_OBJ, 3))
target_bboxes_semcls = np.zeros((MAX_NUM_OBJ))
point_cloud, choices = pc_util.random_sampling(raw_point_cloud,
self.num_points,
return_choices=True)
#ema_point_cloud = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=False)
ema_point_cloud = point_cloud.copy() # 2021.2.28
instance_labels = instance_labels[choices]
semantic_labels = semantic_labels[choices]
target_bboxes_mask[0:instance_bboxes.shape[0]] = 1
target_bboxes[0:instance_bboxes.shape[0], :] = instance_bboxes[:, 0:6]
# ------------------------------- DATA AUGMENTATION ------------------------------
flip_x_axis = 0
flip_y_axis = 0
flip_x_axis_ema = 0
flip_y_axis_ema = 0
rot_mat = np.identity(3)
scale_ratio = np.ones((1, 3))
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
flip_x_axis = 1
point_cloud[:, 0] = -1 * point_cloud[:, 0]
target_bboxes[:, 0] = -1 * target_bboxes[:, 0]
if np.random.random() > 0.5: # 2021.2.28
# Flipping along the YZ plane for ema
flip_x_axis_ema = 1
ema_point_cloud[:, 0] = -1 * ema_point_cloud[:, 0]
if np.random.random() > 0.5:
# Flipping along the XZ plane
flip_y_axis = 1
point_cloud[:, 1] = -1 * point_cloud[:, 1]
target_bboxes[:, 1] = -1 * target_bboxes[:, 1]
if np.random.random() > 0.5: # 2021.2.28
# Flipping along the XZ plane for ema
flip_y_axis_ema = 1
ema_point_cloud[:, 1] = -1 * ema_point_cloud[:, 1]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random() * np.pi /
18) - np.pi / 36 # -5 ~ +5 degree
rot_mat = pc_util.rotz(rot_angle)
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat)
# Augment point cloud scale: 0.85x-1.15x
scale_ratio = np.random.random() * 0.3 + 0.85
scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0)
point_cloud[:, 0:3] *= scale_ratio
target_bboxes[:, 0:3] *= scale_ratio
target_bboxes[:, 3:6] *= scale_ratio
if self.use_height:
point_cloud[:, -1] *= scale_ratio[0, 0]
# compute votes *AFTER* augmentation
# generate votes
# Note: since there's no map between bbox instance labels and
# pc instance_labels (it had been filtered
# in the data preparation step) we'll compute the instance bbox
# from the points sharing the same instance label.
point_votes = np.zeros([self.num_points, 3])
point_votes_mask = np.zeros(self.num_points)
for i_instance in np.unique(instance_labels):
# find all points belong to that instance
ind = np.where(instance_labels == i_instance)[0]
# find the semantic label
if semantic_labels[ind[0]] in DC.nyu40ids:
x = point_cloud[ind, :3]
center = 0.5 * (x.min(0) + x.max(0))
point_votes[ind, :] = center - x
point_votes_mask[ind] = 1.0
point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical
class_ind = [
np.where(DC.nyu40ids == x)[0][0] for x in instance_bboxes[:, -1]
]
# NOTE: set size class as semantic class. Consider use size2class.
size_classes[0:instance_bboxes.shape[0]] = class_ind
size_residuals[0:instance_bboxes.shape[0], :] = \
target_bboxes[0:instance_bboxes.shape[0], 3:6] - DC.mean_size_arr[class_ind, :]
target_bboxes_semcls[0:instance_bboxes.shape[0]] = class_ind
ret_dict = {}
ret_dict['point_clouds'] = point_cloud.astype(np.float32)
ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3]
ret_dict['heading_class_label'] = angle_classes.astype(np.int64)
ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32)
ret_dict['size_class_label'] = size_classes.astype(np.int64)
ret_dict['size_residual_label'] = size_residuals.astype(np.float32)
ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64)
ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32)
ret_dict['vote_label'] = point_votes.astype(np.float32)
ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['supervised_mask'] = np.array(1).astype(np.int64)
scene_label = np.zeros(DC.num_class)
unique_class_ind = list(set(class_ind))
for ind in unique_class_ind:
scene_label[int(ind)] = 1
ret_dict['scene_label'] = scene_label.astype(np.float32)
ret_dict['ema_point_clouds'] = ema_point_cloud.astype(np.float32)
ret_dict['flip_x_axis'] = np.array(flip_x_axis).astype(np.int64)
ret_dict['flip_y_axis'] = np.array(flip_y_axis).astype(np.int64)
ret_dict['rot_mat'] = rot_mat.astype(np.float32)
ret_dict['scale'] = np.array(scale_ratio).astype(np.float32)
ret_dict['flip_x_axis_ema'] = np.array(flip_x_axis_ema).astype(
np.int64) #2021.2.28
ret_dict['flip_y_axis_ema'] = np.array(flip_y_axis_ema).astype(
np.int64) #2021.2.28
return ret_dict
def __getitem__(self, index: int):
id_scan = self.scene_list[index]
assert (id_scan not in self.error_scan)
id_scan_path = os.path.join(self.data_path, id_scan)
point_cloud = np.load(
os.path.join(id_scan_path, '{}.npy'.format('point_cloud')))
ins_vert = np.load(
os.path.join(id_scan_path, '{}.npy'.format('ins_vert'))).squeeze(1)
ins_bbox = np.load(os.path.join(id_scan_path, '{}.npy'.format('bbox')))
points = point_cloud # (N, 3)
center = ins_bbox[:, 0:3] # (B, 10)
bbox_length = ins_bbox[:, 3:6] # (B, 3)
sem_cls = ins_bbox[:, 6:7] # (B, 1)
symmetry = ins_bbox[:, 7:8] # (B, 1)
K = center.shape[0]
# LABELS
if points.shape[0] > self.num_points:
points, choices = pc_util.random_sampling(points,
self.num_points,
return_choices=True)
ins_vert = ins_vert[choices]
elif points.shape[0] < self.num_points:
print('false data')
target_bboxes = np.zeros((MAX_NUM_OBJ, 6), dtype=np.float32)
target_bboxes_mask = np.zeros((MAX_NUM_OBJ), dtype=np.int64)
# target_center = np.zeros((MAX_NUM_OBJ, 3), dtype=np.float32)
# target_rot_q = np.zeros((MAX_NUM_OBJ, 4), dtype=np.float32)
# target_rot_6d = np.zeros((MAX_NUM_OBJ, 6), dtype=np.float32)
# target_scale = np.zeros((MAX_NUM_OBJ, 3), dtype=np.float32)
target_sem_cls = np.zeros((MAX_NUM_OBJ, ), dtype=np.int64)
target_sym = np.zeros((MAX_NUM_OBJ, ), dtype=np.int64)
target_size_classes = np.zeros((MAX_NUM_OBJ, ))
target_size_residuals = np.zeros((MAX_NUM_OBJ, 3))
# target_center[:K] = center[:,0:3]
# target_rot_q[:K] = alignments[:,3:7]
# for k in range(K):
# target_rot_6d[k] = from_q_to_6d(alignments[k,3:7])
# target_scale[:K] = alignments[:,7:10]
target_sem_cls[:K] = sem_cls.squeeze(1)
target_sym[:K] = symmetry.squeeze(1)
target_bboxes[:K, 0:3] = center
target_bboxes[:K, 3:6] = bbox_length
target_bboxes_mask[:K] = 1
# ------------------------------- DATA AUGMENTATION ------------------------------
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
points[:, 0] = -1 * points[:, 0]
target_bboxes[:, 0] = -1 * target_bboxes[:, 0]
if np.random.random() > 0.5:
# Flipping along the XZ plane
points[:, 1] = -1 * points[:, 1]
target_bboxes[:, 1] = -1 * target_bboxes[:, 1]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random() * np.pi /
18) - np.pi / 36 # -5 ~ +5 degree
rot_mat = pc_util.rotz(rot_angle)
points[:, 0:3] = np.dot(points[:, 0:3], np.transpose(rot_mat))
target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat)
target_center = target_bboxes[:, 0:3]
# ====== GENERATE VOTES ======
# compute votes *AFTER* augmentation
# NOTE: i_ins: (1,B) not (0,B-1)
point_votes = np.zeros([self.num_points, 3])
point_votes_mask = np.zeros(self.num_points)
for i_ins in np.unique(ins_vert):
i_ins -= 1
if target_sem_cls[i_ins] in NOT_CARED_IDS or i_ins < 0:
continue
ind = np.where(ins_vert == i_ins + 1)[0]
x = points[ind, :3]
point_votes[ind, :] = x - target_center[i_ins]
point_votes_mask[ind] = 1.0
point_votes = np.tile(point_votes, (1, 3))
target_size_classes[:K] = target_sem_cls[:K]
target_size_residuals[:K, :3] =\
target_bboxes[:K, 3:6] - DC.mean_size_arr[target_sem_cls[:K], :]
# ====== LABELS ======
label = {}
label['point_clouds'] = points.astype(np.float32)
label['center_label'] = target_center.astype(np.float32)
label['heading_class_label'] = np.zeros(
(MAX_NUM_OBJ, )).astype(np.int64)
label['heading_residual_label'] = np.zeros(
(MAX_NUM_OBJ, )).astype(np.float32)
label['size_class_label'] = target_size_classes.astype(np.int64)
label['size_residual_label'] = target_size_residuals.astype(np.float32)
label['sem_cls_label'] = target_sem_cls.astype(np.int64)
label['box_label_mask'] = target_bboxes_mask.astype(np.float32)
label['vote_label'] = point_votes.astype(np.float32)
label['vote_label_mask'] = point_votes_mask.astype(np.int64)
label['scan_idx'] = np.array(index).astype(np.int64)
return label
def __getitem__(self, idx):
"""
Returns a dict with following keys:
point_clouds: (N,3+C)
scan_idx: int scan index in scan_names list
"""
scan_name = self.scan_names[idx]
mesh_vertices = np.load(
os.path.join(self.data_path, scan_name) + '_pc.npz')['pc'] # Nx6
if not self.use_color:
raw_point_cloud = mesh_vertices[:, 0:3] # do not use color for now
else:
raw_point_cloud = mesh_vertices[:, 0:6]
raw_point_cloud[:, 3:] = (raw_point_cloud[:, 3:] -
MEAN_COLOR_RGB) / 256.0
if self.use_height:
floor_height = np.percentile(raw_point_cloud[:, 2], 0.99)
height = raw_point_cloud[:, 2] - floor_height
raw_point_cloud = np.concatenate(
[raw_point_cloud, np.expand_dims(height, 1)], 1)
ret_dict = {}
ema_point_cloud = pc_util.random_sampling(raw_point_cloud,
self.num_points,
return_choices=False)
ret_dict['ema_point_clouds'] = ema_point_cloud.astype(np.float32)
bboxes = np.load(
os.path.join(self.data_path, scan_name) + '_bbox.npy') # K,8
target_bboxes = np.zeros((MAX_NUM_OBJ, 6))
target_bboxes_mask = np.zeros((MAX_NUM_OBJ))
angle_classes = np.zeros((MAX_NUM_OBJ, ))
angle_residuals = np.zeros((MAX_NUM_OBJ, ))
size_classes = np.zeros((MAX_NUM_OBJ, ))
size_residuals = np.zeros((MAX_NUM_OBJ, 3))
target_bboxes_semcls = np.zeros((MAX_NUM_OBJ))
target_bboxes_mask[0:bboxes.shape[0]] = 1
target_bboxes[0:bboxes.shape[0], :] = bboxes[:, 0:6]
for i in range(bboxes.shape[0]):
bbox = bboxes[i]
semantic_class = bbox[7]
angle_class, angle_residual = DC.angle2class(bbox[6])
# NOTE: The mean size stored in size2class is of full length of box edges,
# while in sunrgbd_data.py data dumping we dumped *half* length l,w,h.. so have to time it by 2 here
box3d_size = bbox[3:6] * 2
size_class, size_residual = DC.size2class(
box3d_size, DC.class2type[semantic_class])
angle_classes[i] = angle_class
angle_residuals[i] = angle_residual
size_classes[i] = size_class
size_residuals[i] = size_residual
target_bboxes_semcls[i] = semantic_class
if self.load_labels:
ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3]
ret_dict['heading_class_label'] = angle_classes.astype(np.int64)
ret_dict['heading_residual_label'] = angle_residuals.astype(
np.float32)
ret_dict['size_class_label'] = size_classes.astype(np.int64)
ret_dict['size_residual_label'] = size_residuals.astype(np.float32)
ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64)
ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32)
point_cloud, choices = pc_util.random_sampling(raw_point_cloud,
self.num_points,
return_choices=True)
flip_x_axis = 0
flip_y_axis = 0
rot_angle = 0
rot_mat = np.identity(3)
scale_ratio = np.ones((1, 3))
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
flip_x_axis = 1
point_cloud[:, 0] = -1 * point_cloud[:, 0]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random() * np.pi /
3) - np.pi / 6 # -30 ~ +30 degree
rot_mat = pc_util.rotz(rot_angle)
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
# Augment point cloud scale: 0.85x-1.15x
scale_ratio = np.random.random() * 0.3 + 0.85
scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0)
point_cloud[:, 0:3] *= scale_ratio
if self.use_height:
point_cloud[:, -1] *= scale_ratio[0, 0]
ret_dict['point_clouds'] = point_cloud.astype(np.float32)
ret_dict['flip_x_axis'] = np.array(flip_x_axis).astype(np.int64)
ret_dict['flip_y_axis'] = np.array(flip_y_axis).astype(np.int64)
ret_dict['rot_mat'] = rot_mat.astype(np.float32)
ret_dict['rot_angle'] = np.array(rot_angle).astype(np.float32)
ret_dict['scale'] = np.array(scale_ratio).astype(np.float32)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['supervised_mask'] = np.array(0).astype(np.int64)
return ret_dict
def __getitem__(self, idx):
"""
Returns a dict with following keys:
point_clouds: (N,3+C)
center_label: (MAX_NUM_OBJ,3) for GT box center XYZ
sem_cls_label: (MAX_NUM_OBJ,) semantic class index
angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1
angle_residual_label: (MAX_NUM_OBJ,)
size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER
size_residual_label: (MAX_NUM_OBJ,3)
box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box
point_votes: (N,3) with votes XYZ
point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB.
scan_idx: int scan index in scan_names list
pcl_color: unused
"""
file = self.files[idx]
scan_name = file.split('.')[0]
h5file = h5py.File(os.path.join(self.data_path, file), 'r')
mesh_vertices = np.array(h5file['point_cloud'], dtype=np.float32)
instance_labels = np.array(h5file['instance'], dtype=np.int32)
semantic_labels = np.array(h5file['semantic'], dtype=np.int32)
instance_bboxes = np.array(h5file['bboxes'], dtype=np.float32)
instance_bboxes = instance_bboxes[:, :8]
h5file.close()
# convert PC to z is up.
mesh_vertices[:, [0, 1, 2]] = mesh_vertices[:, [0, 2, 1]]
# convert annotations to z is up.
instance_bboxes[:, [0, 1, 2]] = instance_bboxes[:, [0, 2, 1]]
instance_bboxes[:, [3, 4, 5]] = instance_bboxes[:, [3, 5, 4]]
if not self.use_color:
point_cloud = mesh_vertices[:, 0:3] # do not use color for now
else:
point_cloud = mesh_vertices[:, 0:6]
point_cloud[:, 3:] = point_cloud[:, 3:] - (MEAN_COLOR_RGB) / 256.0
if self.use_height:
floor_height = np.percentile(point_cloud[:, 2], 0.99)
height = point_cloud[:, 2] - floor_height
point_cloud = np.concatenate(
[point_cloud, np.expand_dims(height, 1)], 1)
# ------------------------------- LABELS ------------------------------
target_bboxes = np.zeros((MAX_NUM_OBJ_FT, 6))
target_bboxes_mask = np.zeros((MAX_NUM_OBJ_FT))
angle_classes = np.zeros((MAX_NUM_OBJ_FT, ))
angle_residuals = np.zeros((MAX_NUM_OBJ_FT, ))
size_classes = np.zeros((MAX_NUM_OBJ_FT, ))
size_residuals = np.zeros((MAX_NUM_OBJ_FT, 3))
target_bboxes_mask[0:instance_bboxes.shape[0]] = 1
target_bboxes[0:instance_bboxes.shape[0], :] = instance_bboxes[:, 0:6]
# ------------------------------- DATA AUGMENTATION ------------------------------
if self.augment:
if np.random.random() > 0.5:
# Flipping along the YZ plane
point_cloud[:, 0] = -1 * point_cloud[:, 0]
target_bboxes[:, 0] = -1 * target_bboxes[:, 0]
if np.random.random() > 0.5:
# Flipping along the XZ plane
point_cloud[:, 1] = -1 * point_cloud[:, 1]
target_bboxes[:, 1] = -1 * target_bboxes[:, 1]
# Rotation along up-axis/Z-axis
rot_angle = (np.random.random() * np.pi /
18) - np.pi / 36 # -5 ~ +5 degree
rot_mat = pc_util.rotz(rot_angle)
point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3],
np.transpose(rot_mat))
target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat)
# compute votes *AFTER* augmentation
# generate votes
# Note: since there's no map between bbox instance labels and
# pc instance_labels (it had been filtered
# in the data preparation step) we'll compute the instance bbox
# from the points sharing the same instance label.
num_points = mesh_vertices.shape[0]
point_votes = np.zeros([num_points, 3])
point_votes_mask = np.zeros(num_points)
#DEBUG:
votes_rgb = np.zeros([num_points, 3])
for i_instance in np.unique(instance_labels):
# ignore points not associated with a box
#if i_instance not in instance_bboxes_instance_labels: continue
# find all points belong to that instance
ind = np.where(instance_labels == i_instance)[0]
# find the semantic label
#TODO: change classe labels
if not (semantic_labels[ind[0]] == -1):
x = point_cloud[ind, :3]
center = 0.5 * (x.min(0) + x.max(0))
point_votes[ind, :] = center - x
point_votes_mask[ind] = 1.0
#DEBUG
debug_color = np.random.uniform(0, 1, 3)
votes_rgb[ind, :] = debug_color
point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical
# NOTE: set size class as semantic class. Consider use size2class.
size_classes[0:instance_bboxes.shape[0]] = instance_bboxes[:, -1]
instance_bboxes_sids = instance_bboxes[:, -1]
instance_bboxes_sids = instance_bboxes_sids.astype(np.int)
size_residuals[0:instance_bboxes.shape[0], :] = \
target_bboxes[0:instance_bboxes.shape[0], 3:6] - DC.mean_size_arr[instance_bboxes_sids,:]
#TODO: update angle_classes + residuals
angle_residuals[0:instance_bboxes.shape[0]] = instance_bboxes[:, 6]
#DEBUG
#DEBUG
#DEBUG
#DEBUG
#DEBUG
# DEBUG Point cloud
# DEBUG: viz PC with colors associated to votes
# -- -- debug_pcd = o3d.geometry.PointCloud()
# -- -- debug_point_cloud = point_cloud[:,:3].astype(np.float32)
# -- -- debug_pcd.points = o3d.utility.Vector3dVector(debug_point_cloud)
# -- -- debug_pcd.colors = o3d.utility.Vector3dVector(votes_rgb)
# -- -- o3d.io.write_point_cloud('debug_dataloader_pc.ply',
# -- -- debug_pcd)
# -- -- min_bound = debug_pcd.get_min_bound()
# -- -- max_bound = debug_pcd.get_max_bound()
# -- -- print(min_bound)
# -- -- print(max_bound)
# -- -- # Dump GT bounding boxes
# -- -- gt_center = target_bboxes.astype(np.float32)[:,0:3] # (B,MAX_NUM_OBJ_FT,3)
# -- -- gt_mask = target_bboxes_mask.astype(np.float32) # B,K2
# -- -- gt_heading_class = angle_classes.astype(np.int64) # B,K2
# -- -- gt_heading_residual = angle_residuals.astype(np.float32) # B,K2
# -- -- gt_size_class = size_classes.astype(np.int64) # B,K2
# -- -- gt_size_residual = size_residuals.astype(np.float32) # B,K2,3
# -- -- obbs = []
# -- -- for j in range(gt_center.shape[0]):
# -- -- if gt_mask[j] == 0: continue
# -- -- obb = DC.param2obb(gt_center[j,0:3],
# -- -- gt_heading_class[j],
# -- -- gt_heading_residual[j],
# -- -- gt_size_class[j],
# -- -- gt_size_residual[j])
# -- -- obbs.append(obb)
# -- -- print('stored bboxes: ', len(obbs))
# -- -- if len(obbs)>0:
# -- -- obbs = np.vstack(tuple(obbs)) # (num_gt_objects, 7)
# -- -- pc_util.write_oriented_bbox(obbs, 'debug_dataloader_boxes.ply')
# -- -- stop
#DEBUG
#DEBUG
#DEBUG
#DEBUG
#DEBUG
#DEBUG
#DEBUG
ret_dict = {}
ret_dict['point_clouds'] = point_cloud.astype(np.float32)
ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3]
ret_dict['heading_class_label'] = angle_classes.astype(np.int64)
ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32)
ret_dict['size_class_label'] = size_classes.astype(np.int64)
ret_dict['size_residual_label'] = size_residuals.astype(np.float32)
target_bboxes_semcls = np.zeros((MAX_NUM_OBJ_FT))
target_bboxes_semcls[0:instance_bboxes.shape[0]] = instance_bboxes[:,
-1]
ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64)
ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32)
ret_dict['vote_label'] = point_votes.astype(np.float32)
ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['scan_name'] = scan_name
if self.cache_crops_filename:
files_crops = [x for x in self.files_crops if scan_name in x]
files_crops = [
os.path.join(self.root_data_crops_path, self.split_set,
'votenet_inputs', x) for x in files_crops
]
ret_dict['files_crops'] = files_crops
return ret_dict