wypr/dataset/s3dis/s3dis.py [194:282]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - pcl_color = pcl_color[choices] semantic_labels = semantic_labels[choices] if shape_labels is not None: shape_labels = shape_labels[choices] sem_seg_labels = np.ones_like(semantic_labels) sem_seg_labels = sem_seg_labels * -100 # pytorch default ignore_index: -100 for _c in self.CONFIG.nyu40ids_semseg: sem_seg_labels[semantic_labels == _c] = self.CONFIG.nyu40id2class_semseg[_c] # ------------------------------- 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] target_props[:,0] = -1 * target_props[:,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] target_props[:,1] = -1 * target_props[:,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) target_props = rotate_aligned_boxes(target_props, rot_mat) # ------------- DATA AUGMENTATION for consistency loss ----------------------- point_cloud_aug = copy.deepcopy(point_cloud) target_props_aug = copy.deepcopy(target_props) # Flipping along the YZ plane if np.random.random() > 0.8: point_cloud[:,0] = -1 * point_cloud[:,0] target_props_aug[:,0] = -1 * target_props_aug[:,0] # Flipping along the XZ plane if np.random.random() > 0.8: point_cloud[:,1] = -1 * point_cloud[:,1] target_props_aug[:,1] = -1 * target_props_aug[:,1] # roatation rot_angle = (np.random.random()*np.pi/6) rot_mat = pc_util.rotz(rot_angle) point_cloud_aug[:,0:3] = np.dot(point_cloud_aug[:,0:3], np.transpose(rot_mat)) target_props_aug = rotate_aligned_boxes(target_props_aug, rot_mat) # scale min_s = 0.8; max_s = 2 - min_s # scale = np.random.rand(point_cloud.shape[0]) * (max_s - min_s) + min_s # point_cloud_aug[:, :3] *= scale.reshape(-1, 1) scale = np.random.rand() * (max_s - min_s) + min_s point_cloud_aug[:, :3] *= scale target_props_aug *= scale # jittering jitter_min = 0.95; jitter_max = 2 - jitter_min jitter_scale = np.random.rand(point_cloud.shape[0]) * (jitter_max - jitter_min) + jitter_min point_cloud_aug[:, :3] *= jitter_scale.reshape(-1, 1) # # dropout # num_aug_points_sampled = int(0.9 * point_cloud_aug.shape[0]) # point_cloud_aug, choices_aug = pc_util.random_sampling(point_cloud, num_aug_points_sampled, return_choices=True) # -------------------- RETURN ----------------------- ret_dict = {} ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['gt_boxes'] = target_bboxes.astype(np.float32) ret_dict['gt_boxes_cls'] = target_bboxes_semcls.astype(np.int64) ret_dict['gt_boxes_mask'] = target_bboxes_mask.astype(np.float32) ret_dict['proposals'] = target_props.astype(np.float32) ret_dict['proposal_mask'] = target_props_mask.astype(np.float32) points_rois_idx = pc_util.points_within_boxes(point_cloud, target_props) ret_dict['proposal_points_idx'] = points_rois_idx.astype(np.int64) if self.split_set != "train": ret_dict['num_points'] = num_points ret_dict['point_clouds_all'] = point_cloud_all.astype(np.float32) ret_dict['sem_seg_label_all'] = sem_seg_labels_all.astype(np.int64) ret_dict['shape_labels_all'] = shape_labels_all.astype(np.int64) ret_dict['point_clouds_aug'] = point_cloud_aug.astype(np.float32) ret_dict['proposals_aug'] = target_props_aug.astype(np.float32) ret_dict['sem_seg_label'] = sem_seg_labels ret_dict['shape_labels'] = shape_labels.astype(np.int64) # ret_dict['aug_idx'] = choices_aug.astype(np.int64) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) return ret_dict if __name__=='__main__': pass - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - wypr/dataset/scannet/scannet.py [184:271]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - pcl_color = pcl_color[choices] semantic_labels = semantic_labels[choices] if shape_labels is not None: shape_labels = shape_labels[choices] sem_seg_labels = np.ones_like(semantic_labels) sem_seg_labels = sem_seg_labels * -100 # pytorch default ignore_index: -100 for _c in self.CONFIG.nyu40ids_semseg: sem_seg_labels[semantic_labels == _c] = self.CONFIG.nyu40id2class_semseg[_c] # ------------------------------- 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] target_props[:,0] = -1 * target_props[:,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] target_props[:,1] = -1 * target_props[:,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) target_props = rotate_aligned_boxes(target_props, rot_mat) # ------------- DATA AUGMENTATION for consistency loss ----------------------- point_cloud_aug = copy.deepcopy(point_cloud) target_props_aug = copy.deepcopy(target_props) # Flipping along the YZ plane if np.random.random() > 0.8: point_cloud[:,0] = -1 * point_cloud[:,0] target_props_aug[:,0] = -1 * target_props_aug[:,0] # Flipping along the XZ plane if np.random.random() > 0.8: point_cloud[:,1] = -1 * point_cloud[:,1] target_props_aug[:,1] = -1 * target_props_aug[:,1] # roatation rot_angle = (np.random.random()*np.pi/6) rot_mat = pc_util.rotz(rot_angle) point_cloud_aug[:,0:3] = np.dot(point_cloud_aug[:,0:3], np.transpose(rot_mat)) target_props_aug = rotate_aligned_boxes(target_props_aug, rot_mat) # scale min_s = 0.8; max_s = 2 - min_s # scale = np.random.rand(point_cloud.shape[0]) * (max_s - min_s) + min_s # point_cloud_aug[:, :3] *= scale.reshape(-1, 1) scale = np.random.rand() * (max_s - min_s) + min_s point_cloud_aug[:, :3] *= scale target_props_aug *= scale # jittering jitter_min = 0.95; jitter_max = 2 - jitter_min jitter_scale = np.random.rand(point_cloud.shape[0]) * (jitter_max - jitter_min) + jitter_min point_cloud_aug[:, :3] *= jitter_scale.reshape(-1, 1) # # dropout # num_aug_points_sampled = int(0.9 * point_cloud_aug.shape[0]) # point_cloud_aug, choices_aug = pc_util.random_sampling(point_cloud, num_aug_points_sampled, return_choices=True) # -------------------- RETURN ----------------------- ret_dict = {} ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['gt_boxes'] = target_bboxes.astype(np.float32) ret_dict['gt_boxes_cls'] = target_bboxes_semcls.astype(np.int64) ret_dict['gt_boxes_mask'] = target_bboxes_mask.astype(np.float32) ret_dict['proposals'] = target_props.astype(np.float32) ret_dict['proposal_mask'] = target_props_mask.astype(np.float32) points_rois_idx = pc_util.points_within_boxes(point_cloud, target_props) ret_dict['proposal_points_idx'] = points_rois_idx.astype(np.int64) if self.split_set != "train": ret_dict['num_points'] = num_points ret_dict['point_clouds_all'] = point_cloud_all.astype(np.float32) ret_dict['sem_seg_label_all'] = sem_seg_labels_all.astype(np.int64) ret_dict['shape_labels_all'] = shape_labels_all.astype(np.int64) ret_dict['point_clouds_aug'] = point_cloud_aug.astype(np.float32) ret_dict['proposals_aug'] = target_props_aug.astype(np.float32) ret_dict['sem_seg_label'] = sem_seg_labels ret_dict['shape_labels'] = shape_labels.astype(np.int64) # ret_dict['aug_idx'] = choices_aug.astype(np.int64) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) # ret_dict['pcl_color'] = pcl_color return ret_dict if __name__=='__main__': pass - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -