in datasets/scannet.py [0:0]
def __getitem__(self, idx):
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
pcl_color = point_cloud[:, 3:]
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 ------------------------------
MAX_NUM_OBJ = self.dataset_config.max_num_obj
target_bboxes = np.zeros((MAX_NUM_OBJ, 6), dtype=np.float32)
target_bboxes_mask = np.zeros((MAX_NUM_OBJ), dtype=np.float32)
angle_classes = np.zeros((MAX_NUM_OBJ,), dtype=np.int64)
angle_residuals = np.zeros((MAX_NUM_OBJ,), dtype=np.float32)
raw_sizes = np.zeros((MAX_NUM_OBJ, 3), dtype=np.float32)
raw_angles = np.zeros((MAX_NUM_OBJ,), dtype=np.float32)
if self.augment and self.use_random_cuboid:
(
point_cloud,
instance_bboxes,
per_point_labels,
) = self.random_cuboid_augmentor(
point_cloud, instance_bboxes, [instance_labels, semantic_labels]
)
instance_labels = per_point_labels[0]
semantic_labels = per_point_labels[1]
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]
sem_seg_labels = np.ones_like(semantic_labels) * IGNORE_LABEL
for _c in self.dataset_config.nyu40ids_semseg:
sem_seg_labels[
semantic_labels == _c
] = self.dataset_config.nyu40id2class_semseg[_c]
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 = self.dataset_config.rotate_aligned_boxes(
target_bboxes, rot_mat
)
raw_sizes = target_bboxes[:, 3:6]
point_cloud_dims_min = point_cloud.min(axis=0)
point_cloud_dims_max = point_cloud.max(axis=0)
box_centers = target_bboxes.astype(np.float32)[:, 0:3]
box_centers_normalized = shift_scale_points(
box_centers[None, ...],
src_range=[
point_cloud_dims_min[None, ...],
point_cloud_dims_max[None, ...],
],
dst_range=self.center_normalizing_range,
)
box_centers_normalized = box_centers_normalized.squeeze(0)
box_centers_normalized = box_centers_normalized * target_bboxes_mask[..., None]
mult_factor = point_cloud_dims_max - point_cloud_dims_min
box_sizes_normalized = scale_points(
raw_sizes.astype(np.float32)[None, ...],
mult_factor=1.0 / mult_factor[None, ...],
)
box_sizes_normalized = box_sizes_normalized.squeeze(0)
box_corners = self.dataset_config.box_parametrization_to_corners_np(
box_centers[None, ...],
raw_sizes.astype(np.float32)[None, ...],
raw_angles.astype(np.float32)[None, ...],
)
box_corners = box_corners.squeeze(0)
ret_dict = {}
ret_dict["point_clouds"] = point_cloud.astype(np.float32)
ret_dict["gt_box_corners"] = box_corners.astype(np.float32)
ret_dict["gt_box_centers"] = box_centers.astype(np.float32)
ret_dict["gt_box_centers_normalized"] = box_centers_normalized.astype(
np.float32
)
ret_dict["gt_angle_class_label"] = angle_classes.astype(np.int64)
ret_dict["gt_angle_residual_label"] = angle_residuals.astype(np.float32)
target_bboxes_semcls = np.zeros((MAX_NUM_OBJ))
target_bboxes_semcls[0 : instance_bboxes.shape[0]] = [
self.dataset_config.nyu40id2class[x]
for x in instance_bboxes[:, -1][0 : instance_bboxes.shape[0]]
]
ret_dict["gt_box_sem_cls_label"] = target_bboxes_semcls.astype(np.int64)
ret_dict["gt_box_present"] = target_bboxes_mask.astype(np.float32)
ret_dict["scan_idx"] = np.array(idx).astype(np.int64)
ret_dict["pcl_color"] = pcl_color
ret_dict["gt_box_sizes"] = raw_sizes.astype(np.float32)
ret_dict["gt_box_sizes_normalized"] = box_sizes_normalized.astype(np.float32)
ret_dict["gt_box_angles"] = raw_angles.astype(np.float32)
ret_dict["point_cloud_dims_min"] = point_cloud_dims_min.astype(np.float32)
ret_dict["point_cloud_dims_max"] = point_cloud_dims_max.astype(np.float32)
return ret_dict