def clusters_voxelization(self, clusters_idx, clusters_offset, feats, coords, fullscale, scale, mode): ''' :param clusters_idx: (SumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N, cpu :param clusters_offset: (nCluster + 1), int, cpu :param feats: (N, C), float, cuda :param coords: (N, 3), float, cuda :return: ''' c_idxs = clusters_idx[:, 1].cuda() clusters_feats = feats[c_idxs.long()] clusters_coords = coords[c_idxs.long()] clusters_coords_mean = pointgroup_ops.sec_mean(clusters_coords, clusters_offset.cuda()) # (nCluster, 3), float clusters_coords_mean = torch.index_select(clusters_coords_mean, 0, clusters_idx[:, 0].cuda().long()) # (sumNPoint, 3), float clusters_coords -= clusters_coords_mean clusters_coords_min = pointgroup_ops.sec_min(clusters_coords, clusters_offset.cuda()) # (nCluster, 3), float clusters_coords_max = pointgroup_ops.sec_max(clusters_coords, clusters_offset.cuda()) # (nCluster, 3), float clusters_scale = 1 / ((clusters_coords_max - clusters_coords_min) / fullscale).max(1)[0] - 0.01 # (nCluster), float clusters_scale = torch.clamp(clusters_scale, min=None, max=scale) min_xyz = clusters_coords_min * clusters_scale.unsqueeze(-1) # (nCluster, 3), float max_xyz = clusters_coords_max * clusters_scale.unsqueeze(-1) clusters_scale = torch.index_select(clusters_scale, 0, clusters_idx[:, 0].cuda().long()) clusters_coords = clusters_coords * clusters_scale.unsqueeze(-1) range = max_xyz - min_xyz offset = - min_xyz + torch.clamp(fullscale - range - 0.001, min=0) * torch.rand(3).cuda() + torch.clamp(fullscale - range + 0.001, max=0) * torch.rand(3).cuda() offset = torch.index_select(offset, 0, clusters_idx[:, 0].cuda().long()) clusters_coords += offset assert clusters_coords.shape.numel() == ((clusters_coords >= 0) * (clusters_coords < fullscale)).sum() clusters_coords = clusters_coords.long() clusters_coords = torch.cat([clusters_idx[:, 0].view(-1, 1).long(), clusters_coords.cpu()], 1) # (sumNPoint, 1 + 3) out_coords, inp_map, out_map = pointgroup_ops.voxelization_idx(clusters_coords, int(clusters_idx[-1, 0]) + 1, mode) # output_coords: M * (1 + 3) long # input_map: sumNPoint int # output_map: M * (maxActive + 1) int out_feats = pointgroup_ops.voxelization(clusters_feats, out_map.cuda(), mode) # (M, C), float, cuda spatial_shape = [fullscale] * 3 voxelization_feats = spconv.SparseConvTensor(out_feats, out_coords.int().cuda(), spatial_shape, int(clusters_idx[-1, 0]) + 1) return voxelization_feats, inp_map
def testMerge(self, id): locs = [] locs_float = [] feats = [] batch_offsets = [0] for i, idx in enumerate(id): if self.test_split == 'val': xyz_origin, rgb, label, instance_label = self.test_files[idx] elif self.test_split == 'test': xyz_origin, rgb = self.test_files[idx] else: print("Wrong test split: {}!".format(self.test_split)) exit(0) ### flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append(torch.cat([torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long()], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb)) ### merge all the scenes in the batch batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat(locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(locs, self.batch_size, self.mode) return {'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape}
def testMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] batch_offsets = [0] groups = [] group_to_points = [] for i, idx in enumerate(id): if self.test_split == 'val': xyz_origin, rgb, label, group, point2seg = self.test_files[idx] elif self.test_split == 'test': xyz_origin, rgb = self.test_files[idx] else: print("Wrong test split: {}!".format(self.test_split)) exit(0) xyz_origin = xyz_origin.astype('float32') rgb = rgb.astype('float32') ### flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) point2seg_ori = np.asarray(point2seg) point2seg = point2seg_ori #[valid_idxs] point2seg = point2seg.tolist() group_to_point = {} for j in range(len(point2seg)): if point2seg[j] not in group_to_point: group_to_point[point2seg[j]] = [] group_to_point[point2seg[j]].append(j) for j in range(20): group[j] = list(set(group[j]) & set(group_to_point.keys())) ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb)) labels.append(torch.from_numpy(label)) groups.append(group) group_to_points.append(group_to_point) ### merge all the scenes in the batch batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'groups': groups, 'group2points': group_to_points }
def valMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int batch_offsets = [0] total_inst_num = 0 for i, idx in enumerate(id): xyz_origin, rgb, label = self.val_files[idx] xyz_origin = xyz_origin.astype('float32') rgb = rgb.astype('float32') ### flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) ### crop xyz, valid_idxs = self.crop(xyz) xyz_middle = xyz_middle[valid_idxs] xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] '''instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) ### get instance information inst_num, inst_infos = self.getInstanceInfo(xyz_middle, instance_label.astype(np.int32)) inst_info = inst_infos["instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list instance_label[np.where(instance_label != -100)] += total_inst_num total_inst_num += inst_num''' ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb)) labels.append(torch.from_numpy(label)) '''instance_labels.append(torch.from_numpy(instance_label)) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum)''' ### merge all the scenes in the batch batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) '''instance_labels = torch.cat(instance_labels, 0).long() # long (N) instance_infos = torch.cat(instance_infos, 0).to(torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst)''' spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, #'instance_labels': instance_labels, #'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape }
def trainMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] groups = [] group_to_points = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int batch_offsets = [0] total_inst_num = 0 for i, idx in enumerate(id): xyz_origin, rgb, label, group, point2seg = self.train_files[idx] xyz_origin = xyz_origin.astype('float32') rgb = rgb.astype('float32') ### jitter / flip x / rotation xyz_middle = self.dataAugment(xyz_origin, True, True, True) ### scale xyz = xyz_middle * self.scale ### elastic xyz = self.elastic(xyz, 6 * self.scale // 50, 40 * self.scale / 50) xyz = self.elastic(xyz, 20 * self.scale // 50, 160 * self.scale / 50) ### offset xyz -= xyz.min(0) ### crop xyz, valid_idxs = self.crop(xyz) point2seg_ori = np.asarray(point2seg) point2seg = point2seg_ori[valid_idxs] point2seg = point2seg.tolist() group_to_point = {} for j in range(len(point2seg)): if point2seg[j] not in group_to_point: group_to_point[point2seg[j]] = [] group_to_point[point2seg[j]].append(j) for j in range(20): group[j] = list(set(group[j]) & set(group_to_point.keys())) #print ('keys',group_to_point.keys()) #print ('groups',group) '''np1=np.asarray(group_to_point.keys()) np2=np.asarray(group) #print ('111111111111111111',group_to_point.keys()) #print('222222222222222222',group) g=[] for j in range(20): g.extend(group[j]) #for j in group_to_point: # print (j) # assert (j in g) for j in g: if(j not in group_to_point): print ('j',j,'group',group_to_point.keys(),self.train_file_names[idx],(j in point2seg),(272071 in point2seg),(j in point2seg_ori)) print ('xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx') exit() print (len(g),len(group_to_point))''' #print ('bef',xyz_middle.shape) xyz_middle = xyz_middle[valid_idxs] #print ('aft',xyz_middle.shape) xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] #instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) ### get instance information '''inst_num, inst_infos = self.getInstanceInfo(xyz_middle, instance_label.astype(np.int32)) inst_info = inst_infos["instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list instance_label[np.where(instance_label != -100)] += total_inst_num total_inst_num += inst_num''' ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) groups.append(group) group_to_points.append(group_to_point) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb) + torch.randn(3) * 0.1) labels.append(torch.from_numpy(label)) '''instance_labels.append(torch.from_numpy(instance_label)) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum)''' ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) '''instance_labels = torch.cat(instance_labels, 0).long() # long (N) instance_infos = torch.cat(instance_infos, 0).to(torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst)''' spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, #'instance_labels': instance_labels, #'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'groups': groups, 'group2points': group_to_points }
def Merge(tbl): seg_coords = [] seg_features = [] seg_labels = [] complet_coords = [] complet_invalid = [] voxel_centers = [] complet_invoxel_features = [] complet_labels = [] filenames = [] offset = 0 input_vx = [] stats = [] for idx, example in enumerate(tbl): filename, completion_collection, aliment_collection, segmentation_collection = example '''File Name''' filenames.append(filename) '''Segmentation''' seg_coord = segmentation_collection['coords'] seg_coords.append( torch.cat([ torch.LongTensor(seg_coord.shape[0], 1).fill_(idx), seg_coord.long() ], 1)) seg_labels.append(segmentation_collection['label']) seg_features.append(segmentation_collection['feature']) '''Completion''' complet_coord = aliment_collection['coords'] complet_coords.append( torch.cat([ torch.Tensor(complet_coord.shape[0], 1).fill_(idx), torch.Tensor(complet_coord) ], 1)) input_vx.append(completion_collection['input']) complet_labels.append(completion_collection['label']) complet_invalid.append(completion_collection['invalid']) stats.append(completion_collection['stat']) voxel_centers.append(torch.Tensor(aliment_collection['voxel_centers'])) complet_invoxel_feature = aliment_collection['voxels'] complet_invoxel_feature[:, :, -1] += offset # voxel-to-point mapping in the last column offset += seg_coord.shape[0] complet_invoxel_features.append(torch.Tensor(complet_invoxel_feature)) seg_coords = torch.cat(seg_coords, 0) voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( seg_coords.contiguous(), args['TRAIN']['batch_size'], args['Segmentation']['mode']) spatial_shape = np.clip((seg_coords.max(0)[0][1:] + 1).numpy(), args['Segmentation']['full_scale'][0], None) # long (3) seg_inputs = { 'seg_coords': seg_coords, 'seg_labels': torch.cat(seg_labels, 0), 'seg_features': torch.cat(seg_features, 0), 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'spatial_shape': spatial_shape, } complet_inputs = { 'complet_coords': torch.cat(complet_coords, 0), 'complet_input': torch.cat(input_vx, 0), 'voxel_centers': torch.cat(voxel_centers, 0), 'complet_invalid': torch.cat(complet_invalid, 0), 'complet_labels': torch.cat(complet_labels, 0), 'state': stats, 'complet_invoxel_features': torch.cat(complet_invoxel_features, 0) } return seg_inputs, complet_inputs, filenames
def test_collate_fn(self, id): """ :param id: a list of integers """ locs = [] locs_float = [] feats = [] valid_masks = [] label_abs_paths = [] batch_offsets = [0] for i, idx in enumerate(id): # Load data label_filename = self.label_filenames[idx].strip() rgb_filename = label_filename.replace('segmentation', 'rgb').replace( 'png', 'jpeg') depth_filename = label_filename.replace('segmentation', 'depth') xyz, _, rgb, _, _, valid_mask = \ self.load_data(rgb_filename, depth_filename, label_filename) label_abs_path = '/'.join( label_filename.split('/')[-2:]) # Used for evaluation ### scale xyz_scaled = xyz * self.scale ### offset and crop xyz_scaled -= xyz_scaled.min(0) xyz_scaled, valid_idxs = self.crop(xyz_scaled) # NOTE: for TOD, this doesn't actually crop since I set cfg.max_npoint > #pixels xyz_scaled = xyz_scaled[valid_idxs] xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] if not np.all(valid_idxs): raise Exception("valid_idxs is not all True...") raise Exception("valid_idxs is not all True...") ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz_scaled.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz_scaled.shape[0], 1).fill_(i), torch.from_numpy(xyz_scaled).long() ], 1)) # Used for voxelization locs_float.append( torch.from_numpy(xyz)) # XYZ coords used in the network feats.append( torch.from_numpy(rgb) ) # Can set this to be torch.zeros((N,0), dtype=torch.float) if I don't want RGB valid_masks.append(torch.from_numpy(valid_mask)) label_abs_paths.append(label_abs_path) ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) valid_mask = torch.cat(valid_masks, 0).bool() # bool (N) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'valid_mask': valid_mask, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'label_abs_path': label_abs_paths }
def train_collate_fn(self, id): """ :param id: a list of integers """ locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] valid_masks = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int batch_offsets = [0] total_inst_num = 0 for i, idx in enumerate(id): # Load data label_filename = self.label_filenames[idx].strip() rgb_filename = label_filename.replace('segmentation', 'rgb').replace( 'png', 'jpeg') depth_filename = label_filename.replace('segmentation', 'depth') xyz_origin, xyz_aug, rgb, label, instance_label, valid_mask = \ self.load_data(rgb_filename, depth_filename,label_filename) ### scale xyz = xyz_aug * self.scale ### offset and crop xyz -= xyz.min(0) xyz, valid_idxs = self.crop(xyz) # NOTE: for TOD, this doesn't actually crop since I set cfg.max_npoint > #pixels xyz = xyz[valid_idxs] xyz_aug = xyz_aug[valid_idxs] xyz_origin = xyz_origin[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) if not np.all(valid_idxs): raise Exception("valid_idxs is not all True...") ### get instance information inst_num, inst_infos = self.getInstanceInfo( xyz_origin, instance_label.astype(np.int32)) inst_info = inst_infos[ "instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list instance_label[ instance_label != cfg.ignore_label] += total_inst_num total_inst_num += inst_num ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) # Used for voxelization locs_float.append( torch.from_numpy(xyz_aug)) # XYZ coords used in the network feats.append( torch.from_numpy(rgb) ) # Can set this to be torch.zeros((N,0), dtype=torch.float) if I don't want RGB labels.append(torch.from_numpy(label)) instance_labels.append(torch.from_numpy(instance_label)) valid_masks.append(torch.from_numpy(valid_mask)) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum) ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) instance_labels = torch.cat(instance_labels, 0).long() # long (N) valid_mask = torch.cat(valid_masks, 0).bool() # bool (N) instance_infos = torch.cat(instance_infos, 0).to( torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, 'valid_mask': valid_mask, 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape }
def collate_fn(self, id): """ :param id: a list of integers """ locs = [] locs_float = [] feats = [] valid_masks = [] label_abs_paths = [] batch_offsets = [0] for i, idx in enumerate(id): # Load data pcd_filename = self.pcd_files[idx] temp_idx = pcd_filename.split('/').index( 'OCID-dataset' ) # parse something like this: /data/OCID-dataset/YCB10/table/top/curved/seq36/pcd/result_2018-08-24-15-13-13.pcd label_abs_path = '/'.join(pcd_filename.split('/')[temp_idx + 1:]) xyz, rgb, valid_mask = self.load_data(pcd_filename) ### scale xyz_scaled = xyz * self.scale ### offset and crop xyz_scaled -= xyz_scaled.min(0) xyz_scaled, valid_idxs = self.crop(xyz_scaled) # NOTE: for TOD, this doesn't actually crop since I set cfg.max_npoint > #pixels xyz_scaled = xyz_scaled[valid_idxs] xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] if not np.all(valid_idxs): raise Exception("valid_idxs is not all True...") ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz_scaled.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz_scaled.shape[0], 1).fill_(i), torch.from_numpy(xyz_scaled).long() ], 1)) # Used for voxelization locs_float.append( torch.from_numpy(xyz)) # XYZ coords used in the network feats.append( torch.from_numpy(rgb) ) # Can set this to be torch.zeros((N,0), dtype=torch.float) if I don't want RGB valid_masks.append(torch.from_numpy(valid_mask)) label_abs_paths.append(label_abs_path) ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) valid_mask = torch.cat(valid_masks, 0).bool() # bool (N) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'valid_mask': valid_mask, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'label_abs_path': label_abs_paths }
def testMerge(self, id): ret_dict = {} # variables for backbone point_locs = [] # (N, 4) (sample_index, xyz) point_coords = [] # (N, 6) (shifted_xyz, original_xyz) point_feats = [] # (N, 3) (rgb) point_semantic_labels = [] point_instance_infos = [] point_instance_labels = [] # (N) batch_offsets = [0] total_inst_num = 0 for i, idx in enumerate(id): assert self.test_split in ['val', 'test'] if not self.cache: if self.test_split == 'val': xyz_origin, rgb, label, instance_label = self.test_files[idx] elif self.test_split == 'test': xyz_origin, rgb = self.test_files[idx] else: fn = self.test_file_names[idx].split('/')[-1][:12] xyz_origin = SA.attach("shm://{}_xyz".format(fn)).copy() rgb = SA.attach("shm://{}_rgb".format(fn)).copy() ### jitter / flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) # variables for backbone point_locs.append(torch.cat((torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long()), dim=1)) # (N, 4) (sample_index, xyz) point_coords.append(torch.from_numpy(np.concatenate((xyz_middle, xyz_origin), axis=1))) # (N, 6) (shifted_xyz, original_xyz) point_feats.append(torch.from_numpy(rgb) + torch.randn(3) * 0.1) # (N, 3) (rgb) if self.test_split == 'val': point_instance_labels.append(torch.from_numpy(instance_label)) # (N) if self.test_split == 'val': point_semantic_labels.append(torch.from_numpy(label)) # (N) xyz, valid_idxs = self.crop(xyz) instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) ### get instance information inst_num, inst_infos = self.getInstanceInfo(xyz_middle, instance_label.astype(np.int32), label.copy()) instance_infos = inst_infos["instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) point_instance_infos.append(torch.from_numpy(instance_infos)) # (N, 9) (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) point_locs = torch.cat(point_locs, 0) # (N) (sample_index) point_coords = torch.cat(point_coords, 0).to(torch.float32) # (N, 6) (shifted_xyz, original_xyz) point_feats = torch.cat(point_feats, 0) # (N, 3) (rgb) if self.test_split == 'val': point_instance_labels = torch.cat(point_instance_labels, 0).long() # (N) # variables for backbone ret_dict['point_locs'] = point_locs # (N, 4) (sample_index, xyz) ret_dict['point_coords'] = point_coords # (N, 6) (shifted_xyz, original_xyz) ret_dict['point_feats'] = point_feats # (N, 3) (rgb) spatial_shape = np.clip((point_locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(point_locs, self.batch_size, self.mode) if ('occupancy' in self.model_mode.split('_')) and (self.test_split == 'val'): voxel_instance_labels = point_instance_labels[v2p_map[:, 1].long()] ignore_voxel_index = (voxel_instance_labels == -100) voxel_instance_labels[ignore_voxel_index] = voxel_instance_labels.max() + 1 num_instance_occupancy = torch_scatter.scatter_add( torch.ones_like(voxel_instance_labels), voxel_instance_labels ) voxel_occupancy_labels = num_instance_occupancy[voxel_instance_labels] voxel_occupancy_labels[ignore_voxel_index] = 0 voxel_instance_labels[ignore_voxel_index] = -100 ret_dict['voxel_instance_labels'] = voxel_instance_labels ret_dict['voxel_occupancy_labels'] = voxel_occupancy_labels # variables for point-wise predictions ret_dict['voxel_locs'] = voxel_locs # (nVoxel, 4) ret_dict['p2v_map'] = p2v_map # (N) ret_dict['v2p_map'] = v2p_map # (nVoxel, 19?) # variables for other uses ret_dict['id'] = id ret_dict['batch_offsets'] = batch_offsets # int (B+1) ret_dict['spatial_shape'] = spatial_shape # long (3) return ret_dict
def valMerge(self, id): ret_dict = {} # variables for backbone point_locs = [] # (N, 4) (sample_index, xyz) point_coords = [] # (N, 6) (shifted_xyz, original_xyz) point_feats = [] # (N, 3) (rgb) # variables for point-wise predictions point_semantic_labels = [] # (N) point_instance_labels = [] # (N) point_instance_infos = [] # (N, 9) instance_centers = [] # (nInst, 3) (instance_xyz) instance_sizes = [] # instance_pointnum = [] # (total_nInst), int batch_offsets = [0] total_inst_num = 0 for i, idx in enumerate(id): if self.cache: fn = self.val_file_names[idx].split('/')[-1][:12] xyz_origin = SA.attach("shm://{}_xyz".format(fn)).copy() rgb = SA.attach("shm://{}_rgb".format(fn)).copy() label = SA.attach("shm://{}_label".format(fn)).copy() instance_label = SA.attach("shm://{}_instance_label".format(fn)).copy() else: xyz_origin, rgb, label, instance_label = self.val_files[idx] ### jitter / flip x / rotation xyz_middle = self.dataAugment(xyz_origin, True, True, True) # xyz_middle = self.dataAugment(xyz_origin, False, False, False) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) ### crop xyz, valid_idxs = self.crop(xyz) xyz_origin = xyz_origin[valid_idxs] xyz_middle = xyz_middle[valid_idxs] xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) ### get instance information inst_num, inst_infos = self.getInstanceInfo(xyz_middle, instance_label.astype(np.int32), label.copy()) instance_infos = inst_infos["instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list inst_centers = inst_infos['instance_center'] # (nInst, 3) (cx, cy, cz) inst_size = inst_infos['instance_size'] inst_label = inst_infos['instance_label'] instance_label[np.where(instance_label != -100)] += total_inst_num total_inst_num += inst_num ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) # variables for backbone point_locs.append(torch.cat((torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long()), dim=1)) # (N, 4) (sample_index, xyz) point_coords.append(torch.from_numpy(np.concatenate((xyz_middle, xyz_origin), axis=1))) # (N, 6) (shifted_xyz, original_xyz) point_feats.append(torch.from_numpy(rgb) + torch.randn(3) * 0.1) # (N, 3) (rgb) # variables for point-wise predictions point_semantic_labels.append(torch.from_numpy(label)) # (N) point_instance_labels.append(torch.from_numpy(instance_label)) # (N) point_instance_infos.append(torch.from_numpy(instance_infos)) # (N, 9) (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) instance_centers.append( torch.cat( (torch.DoubleTensor(len(inst_centers), 1).fill_(i), torch.from_numpy(np.array(inst_centers))), dim=1 ) ) # (nInst, 4) (sample_index, instance_center_xyz) instance_sizes.append( torch.cat( (torch.DoubleTensor(torch.tensor(inst_size).shape[0], 1).fill_(i), torch.tensor(inst_size)), dim=1 ) ) # variable for other uses instance_pointnum.extend(inst_pointnum) ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) point_locs = torch.cat(point_locs, 0) # (N) (sample_index) point_coords = torch.cat(point_coords, 0).to(torch.float32) # (N, 6) (shifted_xyz, original_xyz) point_feats = torch.cat(point_feats, 0) # (N, 3) (rgb) # variables for point-wise predictions point_semantic_labels = torch.cat(point_semantic_labels, 0).long() # (N) point_instance_labels = torch.cat(point_instance_labels, 0).long() # (N) point_instance_infos = torch.cat(point_instance_infos, 0).to(torch.float32) # (N, 9) (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) instance_centers = torch.cat(instance_centers, 0) # (nInst, 3) (instance_center_xyz) instance_sizes = torch.cat(instance_sizes, dim=0) # variables for backbone ret_dict['point_locs'] = point_locs # (N, 4) (sample_index, xyz) ret_dict['point_coords'] = point_coords # (N, 6) (shifted_xyz, original_xyz) ret_dict['point_feats'] = point_feats # (N, 3) (rgb) ret_dict['point_semantic_labels'] = point_semantic_labels # (N) ret_dict['point_instance_labels'] = point_instance_labels # (N) ret_dict['point_instance_infos'] = point_instance_infos # (N, 9) ret_dict['instance_centers'] = instance_centers # (nInst, 3) (instance_xyz) ret_dict['instance_sizes'] = instance_sizes # variable for other uses instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst) spatial_shape = np.clip((point_locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(point_locs, self.batch_size, self.mode) if 'occupancy' in self.model_mode.split('_'): voxel_instance_labels = point_instance_labels[v2p_map[:, 1].long()] ignore_voxel_index = (voxel_instance_labels == -100) voxel_instance_labels[ignore_voxel_index] = voxel_instance_labels.max() + 1 num_instance_occupancy = torch_scatter.scatter_add( torch.ones_like(voxel_instance_labels), voxel_instance_labels ) voxel_occupancy_labels = num_instance_occupancy[voxel_instance_labels] voxel_occupancy_labels[ignore_voxel_index] = 0 voxel_instance_labels[ignore_voxel_index] = -100 ret_dict['voxel_instance_labels'] = voxel_instance_labels ret_dict['voxel_occupancy_labels'] = voxel_occupancy_labels if self.voxel_center_prediction['activate']: voxel_coords = point_coords[v2p_map[:, 1].long(), :3] voxel_center_probs = [] for batch_idx in range(self.batch_size): voxel_batch_idx = (voxel_locs[:, 0] == batch_idx) voxel_center_probability = generate_adaptive_heatmap( voxel_coords[voxel_batch_idx].double(), instance_centers[instance_centers[:, 0] == batch_idx, 1:], instance_sizes[instance_sizes[:, 0] == batch_idx, 1:], min_IoU=self.min_IoU, )['heatmap'] voxel_center_probs.append(voxel_center_probability) voxel_center_probs_labels = torch.cat(voxel_center_probs).to(torch.float32) voxel_instance_info = point_instance_infos[v2p_map[:, 1].long(), :] voxel_center_offset_labels = voxel_instance_info[:, :3] - voxel_coords voxel_center_semantic_labels = point_semantic_labels[v2p_map[:, 1].long()] voxel_center_instance_labels = point_instance_labels[v2p_map[:, 1].long()] ret_dict['voxel_center_probs_labels'] = voxel_center_probs_labels ret_dict['voxel_center_offset_labels'] = voxel_center_offset_labels ret_dict['voxel_center_semantic_labels'] = voxel_center_semantic_labels ret_dict['voxel_center_instance_labels'] = voxel_center_instance_labels # variables for point-wise predictions ret_dict['voxel_locs'] = voxel_locs # (nVoxel, 4) ret_dict['p2v_map'] = p2v_map # (N) ret_dict['v2p_map'] = v2p_map # (nVoxel, 19?) # variables for other uses ret_dict['instance_pointnum'] = instance_pointnum # (nInst) # currently used in Jiang_PointGroup ret_dict['id'] = id ret_dict['batch_offsets'] = batch_offsets # int (B+1) ret_dict['spatial_shape'] = spatial_shape # long (3) return ret_dict
def trainMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int batch_instance_offsets = [] batch_offsets = [0] result = { 'lang_feat': [], 'lang_len': [], 'object_id': [], 'ann_id': [], 'object_cat': [], 'load_time': [], } total_inst_num = 0 for i, idx in enumerate(id): xyz_origin = id[i]['cords'].copy() rgb = id[i]['colors'].copy() label = id[i]['labels'].copy() instance_label = id[i]['instance_labels'].copy() object_id = id[i]['object_id'].copy() scene_id = id[i]['scene_id'].copy() #print('scene_id: ' + str(scene_id)) instance_num = int(instance_label.max()) if instance_num < object_id: print('Number of instances less than object_id') for key in result: result[key].append(id[i][key]) ### jitter / flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * cfg.scale ### elastic #xyz = self.elastic(xyz, 6 * cfg.scale // 50, 40 * cfg.scale / 50) #xyz = self.elastic(xyz, 20 * cfg.scale // 50, 160 * cfg.scale / 50) ### offset xyz -= xyz.min(0) ### crop #xyz, valid_idxs = self.crop(xyz) #xyz_middle = xyz_middle[valid_idxs] #xyz = xyz[valid_idxs] #rgb = rgb[valid_idxs] #label = label[valid_idxs] #instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) ### get instance information inst_num, inst_infos = self.getInstanceInfo( xyz_middle, instance_label.astype(np.int32), object_id) inst_info = inst_infos[ "instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list batch_instance_offsets.append(total_inst_num) instance_label[np.where(instance_label != -100)] += total_inst_num total_inst_num += inst_num ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb) + torch.randn(3) * 0.1) labels.append(torch.from_numpy(label)) instance_labels.append( torch.from_numpy(instance_label.astype(np.int32))) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum) ### merge all the scenes in the batchd batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) instance_labels = torch.cat(instance_labels, 0).long() # long (N) instance_infos = torch.cat(instance_infos, 0).to( torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), cfg.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, cfg.batch_size, cfg.mode) for key in result: result[key] = torch.from_numpy(np.array(result[key])) result['locs'] = locs result['voxel_locs'] = voxel_locs result['p2v_map'] = p2v_map result['v2p_map'] = v2p_map result['locs_float'] = locs_float result['labels'] = labels result['instance_labels'] = instance_labels result['instance_info'] = instance_infos result['instance_pointnum'] = instance_pointnum result['id'] = id result['offsets'] = batch_offsets result['spatial_shape'] = spatial_shape result['feats'] = feats result['batch_instance_offsets'] = np.array(batch_instance_offsets) return result
def testMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] batch_offsets = [0] groups = [] group_to_points = [] group_fulls = [] group_to_point_fulls = [] for i, idx in enumerate(id): if 1: xyz_origin, rgb, label, group, point2seg, group_full, point2seg_full, name = self.test_files[ idx] unary = torch.from_numpy( np.load(self.data_root + '/unary_pred/' + name.split('/')[-1][:12] + '.npy')) unary_feat = torch.from_numpy( np.load(self.data_root + '/unary_feat/' + name.split('/')[-1][:12] + '.npy')) pairwise = torch.from_numpy( np.load(self.data_root + '/rel_feat/' + name.split('/')[-1][:12] + '.npy')) prod = torch.from_numpy( np.load(self.data_root + '/rel_pred/' + name.split('/')[-1][:12] + '.npy')) xyz_origin = xyz_origin.astype('float32') rgb = rgb.astype('float32') ### flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) point2seg_ori = np.asarray(point2seg) point2seg = point2seg_ori #[valid_idxs] point2seg = point2seg.tolist() group_to_point = {} for j in range(len(point2seg)): if point2seg[j] not in group_to_point: group_to_point[point2seg[j]] = [] group_to_point[point2seg[j]].append(j) for j in range(20): group[j] = list(set(group[j]) & set(group_to_point.keys())) point2seg_ori_full = np.asarray(point2seg_full) point2seg_full = point2seg_ori_full #[valid_idxs] point2seg_full = point2seg_full.tolist() group_to_point_full = {} for j in range(len(point2seg_full)): if point2seg_full[j] not in group_to_point_full: group_to_point_full[point2seg_full[j]] = [] group_to_point_full[point2seg_full[j]].append(j) group_full = list( set(group_full) & set(group_to_point_full.keys())) ### merge the scene to the batch batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb)) labels.append(torch.from_numpy(label)) groups.append(group) group_fulls.append(group_full) group_to_points.append(group_to_point) group_to_point_fulls.append(group_to_point_full) ### merge all the scenes in the batch batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) #print ('feats',feats,feats.shape) return { 'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'groups': groups, 'group2points': group_to_points, 'group_fulls': group_fulls, 'group2point_fulls': group_to_point_fulls, 'unary': unary, 'pairwise': pairwise, 'prod': prod, 'unary_feat': unary_feat }
def valMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int lang_feats = [] lang_lens = [] lang_ids = [] ann_ids = [] object_ids = [] object_classes = [] batch_offsets = [0] scene_ids = [] total_inst_num = 0 for i, idx in enumerate(id): #get object scene_id = self.val_data[idx]["scene_id"] object_id = int(self.val_data[idx]["object_id"]) object_name = " ".join(self.val_data[idx]["object_name"].split("_")) ann_id = self.val_data[idx]["ann_id"] #get language features lang_feat = self.lang[scene_id][str(object_id)][ann_id] lang_len = len(self.scanrefer[idx]["token"]) + 2 lang_len = lang_len if lang_len <= cfg.TRAIN_MAX_DES_LEN + 2 else cfg.TRAIN_MAX_DES_LEN + 2 #get scene data data_file = os.path.join(self.data_root,self.dataset,'{}_pointgroup.pth'.format(scene_id)) xyz_origin, rgb, label, instance_label = torch.load(data_file) ### flip x / rotation xyz_middle = self.dataAugment(xyz_origin, False, True, True) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) ### get instance information inst_num, inst_infos, _ = self.getInstanceInfo(xyz_middle, instance_label.astype(np.int32),object_id, eval=True) inst_info = inst_infos["instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list instance_label[np.where(instance_label != -100)] += total_inst_num total_inst_num += inst_num ### merge the scene to the batch scene_ids.append(scene_id) ann_ids.append(int(ann_id)) object_ids.append(int(object_id)) batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append(torch.cat([torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long()], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feat = torch.from_numpy(rgb) if self.use_multiview: multiview = torch.from_numpy(self.multiview_data[scene_id][:]) feat = torch.cat([feat,multiview],1) feats.append(feat) labels.append(torch.from_numpy(label)) instance_labels.append(torch.from_numpy(instance_label)) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum) lang_feats.append(torch.from_numpy(lang_feat).unsqueeze(0)) lang_lens.append(lang_len) lang_ids.append(torch.from_numpy(self.lang_ids[scene_id][str(object_id)][ann_id]).unsqueeze(0)) ### merge all the scenes in the batch ann_ids = torch.tensor(ann_ids, dtype=torch.long) object_ids = torch.tensor(object_ids, dtype=torch.long) batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat(locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) instance_labels = torch.cat(instance_labels, 0).long() # long (N) instance_infos = torch.cat(instance_infos, 0).to(torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst) lang_feats = torch.cat(lang_feats,0).to(torch.float32) lang_lens = torch.tensor(lang_lens,dtype=torch.long) lang_ids = torch.cat(lang_ids,0).to(torch.long) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(locs, self.batch_size, self.mode) return {'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'lang_feat': lang_feats, 'lang_len': lang_lens, 'lang_ids': lang_ids, 'ann_id': ann_ids, 'object_id': object_ids, 'scene_id': scene_ids}
def trainMerge(self, id): locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int target_instance_labels = [] target_instance_pointnum = [] lang_feats = [] lang_lens = [] lang_ids = [] ann_ids = [] object_ids = [] object_classes = [] batch_offsets = [0] total_inst_num = 0 for i, idx in enumerate(id): #get object scene_id = self.train_data[idx]["scene_id"] object_id = int(self.train_data[idx]["object_id"]) object_name = " ".join(self.train_data[idx]["object_name"].split("_")) ann_id = self.train_data[idx]["ann_id"] #get language features lang_feat = self.lang[scene_id][str(object_id)][ann_id] lang_len = len(self.scanrefer[idx]["token"]) + 2 lang_len = lang_len if lang_len <= cfg.TRAIN_MAX_DES_LEN + 2 else cfg.TRAIN_MAX_DES_LEN + 2 #get scene data data_file = os.path.join(self.data_root,self.dataset,'{}_pointgroup.pth'.format(scene_id)) xyz_origin, rgb, label, instance_label = torch.load(data_file) #instance_bboxes = np.load(os.path.join(self.data_root,self.dataset + '_votenet',scene_id)+'_aligned_bbox.npy') ### jitter / flip x / rotation xyz_middle = self.dataAugment(xyz_origin, True, True, True) ### scale xyz = xyz_middle * self.scale ### elastic xyz = self.elastic(xyz, 6 * self.scale // 50, 40 * self.scale / 50) xyz = self.elastic(xyz, 20 * self.scale // 50, 160 * self.scale / 50) ### offset xyz -= xyz.min(0) ### crop xyz, valid_idxs = random_sampling(xyz, self.max_npoint, return_choices=True) xyz_middle = xyz_middle[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] instance_label = self.getCroppedInstLabel(instance_label, valid_idxs) ### get instance information inst_num, inst_infos, target_inst_pointnum = self.getInstanceInfo(xyz_middle, instance_label.astype(np.int32),object_id) inst_info = inst_infos["instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list instance_label[np.where(instance_label != -100)] += total_inst_num #get target object information target_instance_id = object_id + total_inst_num target_instance_label = np.where(instance_label == target_instance_id, instance_label, -100) #only keep captioning target total_inst_num += inst_num ### merge the scene to the batch ann_ids.append(int(ann_id)) object_ids.append(int(object_id)) batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append(torch.cat([torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long()], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feat = torch.from_numpy(rgb) + torch.randn(3) * 0.1 if self.use_multiview: multiview = torch.from_numpy(self.multiview_data[scene_id][:])[valid_idxs] feat = torch.cat([feat,multiview],1) feats.append(feat) labels.append(torch.from_numpy(label)) instance_labels.append(torch.from_numpy(instance_label)) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum) target_instance_labels.append(torch.from_numpy(target_instance_label)) target_instance_pointnum.append(target_inst_pointnum) lang_feats.append(torch.from_numpy(lang_feat).unsqueeze(0)) lang_lens.append(lang_len) lang_ids.append(torch.from_numpy(self.lang_ids[scene_id][str(object_id)][ann_id]).unsqueeze(0)) ### merge all the scenes in the batchd ann_ids = torch.tensor(ann_ids, dtype=torch.long) object_ids = torch.tensor(object_ids, dtype=torch.long) batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat(locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) instance_labels = torch.cat(instance_labels, 0).long() # long (N) instance_infos = torch.cat(instance_infos, 0).to(torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst) target_instance_labels = torch.cat(target_instance_labels,0).long() target_instance_pointnum = torch.tensor(target_instance_pointnum, dtype=torch.int) lang_feats = torch.cat(lang_feats,0).to(torch.float32) lang_lens = torch.tensor(lang_lens,dtype=torch.long) lang_ids = torch.cat(lang_ids,0).to(torch.long) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx(locs, self.batch_size, self.mode) return {'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, 'target_instance_labels': target_instance_labels, 'target_instance_pointnum': target_instance_pointnum, 'id': id, 'offsets': batch_offsets, 'spatial_shape': spatial_shape, 'lang_feat': lang_feats, 'lang_len': lang_lens, 'lang_ids': lang_ids, 'ann_id': ann_ids, 'object_id': object_ids }
def trainMerge(self, id): start = time.time() ## PointGroup Input ## locs = [] locs_float = [] feats = [] labels = [] instance_labels = [] instance_infos = [] # (N, 9) instance_pointnum = [] # (total_nInst), int batch_offsets = [0] total_inst_num = 0 ## ScanRefer Input ## lang_feats = [] lang_lens = [] object_cats = [] object_ids = [] for i, idx in enumerate(id): scene_id = self.scanrefer[idx]["scene_id"] object_id = int(self.scanrefer[idx]["object_id"]) object_name = " ".join( self.scanrefer[idx]["object_name"].split("_")) ann_id = self.scanrefer[idx]["ann_id"] # get language features lang_feat = self.lang[scene_id][str(object_id)][ann_id] lang_len = len(self.scanrefer[idx]["token"]) lang_len = lang_len if lang_len <= CONF.TRAIN.MAX_DES_LEN else CONF.TRAIN.MAX_DES_LEN # get pc mesh_vertices = self.scene_data[scene_id]["mesh_vertices"] instance_label = self.scene_data[scene_id]["instance_labels"] semantic_label = self.scene_data[scene_id]["semantic_labels"] instance_bboxes = self.scene_data[scene_id]["instance_bboxes"] 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:6] = (point_cloud[:, 3:6] - MEAN_COLOR_RGB) / 256.0 pcl_color = point_cloud[:, 3:6] if self.use_normal: normals = mesh_vertices[:, 6:9] point_cloud = np.concatenate([point_cloud, normals], 1) if self.use_multiview: # load multiview database pid = mp.current_process().pid if pid not in self.multiview_data: self.multiview_data[pid] = h5py.File(MULTIVIEW_DATA, "r", libver="latest") multiview = self.multiview_data[pid][scene_id] point_cloud = np.concatenate([point_cloud, multiview], 1) # 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) # Prepare_data_inst.py (PG) point_cloud = np.ascontiguousarray(point_cloud - point_cloud.mean(0)) pcl_color = np.ascontiguousarray(pcl_color[:, 3:6]) / 127.5 - 1 #TODO: Random Sampling point_cloud, choices = random_sampling(point_cloud, self.num_points, return_choices=True) instance_label = instance_label[choices] semantic_label = semantic_label[choices] pcl_color = pcl_color[choices] xyz_origin = point_cloud label = semantic_label rgb = pcl_color if self.data_augmentation: # TODO: Data augmentation ### jitter / flip x / rotation xyz_middle = self.dataAugment(xyz_origin, True, True, True) ### scale xyz = xyz_middle * self.scale ### elastic xyz = self.elastic(xyz, 6 * self.scale // 50, 40 * self.scale / 50) xyz = self.elastic(xyz, 20 * self.scale // 50, 160 * self.scale / 50) ### offset xyz -= xyz.min(0) ### crop xyz, valid_idxs = self.crop(xyz) xyz_middle = xyz_middle[valid_idxs] xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] instance_label = self.getCroppedInstLabel( instance_label, valid_idxs) ### get instance information if not self.data_augmentation: xyz_middle = self.dataAugment(xyz_origin, False, False, False) ### scale xyz = xyz_middle * self.scale ### offset xyz -= xyz.min(0) ### crop xyz, valid_idxs = self.crop(xyz) xyz_middle = xyz_middle[valid_idxs] xyz = xyz[valid_idxs] rgb = rgb[valid_idxs] label = label[valid_idxs] instance_label = self.getCroppedInstLabel( instance_label, valid_idxs) inst_num, inst_infos = self.getInstanceInfo( xyz_middle, instance_label.astype(np.int32)) inst_info = inst_infos[ "instance_info"] # (n, 9), (cx, cy, cz, minx, miny, minz, maxx, maxy, maxz) inst_pointnum = inst_infos["instance_pointnum"] # (nInst), list instance_label[np.where(instance_label != -100)] += total_inst_num object_id += total_inst_num total_inst_num += inst_num ### merge the scene to the batch (PG) batch_offsets.append(batch_offsets[-1] + xyz.shape[0]) locs.append( torch.cat([ torch.LongTensor(xyz.shape[0], 1).fill_(i), torch.from_numpy(xyz).long() ], 1)) locs_float.append(torch.from_numpy(xyz_middle)) feats.append(torch.from_numpy(rgb)) #+ torch.randn(3) * 0.1) labels.append(torch.from_numpy(label)) instance_labels.append( torch.from_numpy(instance_label.astype(np.int64))) instance_infos.append(torch.from_numpy(inst_info)) instance_pointnum.extend(inst_pointnum) ### merge the scene to the batch (SR) # TODO: Check shape of lang_feats, lang_len, object_cats lang_feats.append(torch.from_numpy(lang_feat.astype(np.float32))) lang_lens.append( torch.from_numpy(np.array(lang_len).astype(np.int64))) object_cat = self.raw2label[ object_name] if object_name in self.raw2label else 17 object_cats.append( torch.from_numpy(np.array(object_cat).astype(np.int64))) object_ids.append(object_id) ### merge all the scenes in the batchd (PG) batch_offsets = torch.tensor(batch_offsets, dtype=torch.int) # int (B+1) locs = torch.cat( locs, 0) # long (N, 1 + 3), the batch item idx is put in locs[:, 0] locs_float = torch.cat(locs_float, 0).to(torch.float32) # float (N, 3) feats = torch.cat(feats, 0) # float (N, C) labels = torch.cat(labels, 0).long() # long (N) instance_labels = torch.cat(instance_labels, 0).long() # long (N) instance_infos = torch.cat(instance_infos, 0).to( torch.float32) # float (N, 9) (meanxyz, minxyz, maxxyz) instance_pointnum = torch.tensor(instance_pointnum, dtype=torch.int) # int (total_nInst) spatial_shape = np.clip((locs.max(0)[0][1:] + 1).numpy(), self.full_scale[0], None) # long (3) ### voxelize voxel_locs, p2v_map, v2p_map = pointgroup_ops.voxelization_idx( locs, self.batch_size, self.mode) ### SC lang_feats = torch.cat(lang_feats, 0).reshape( (self.batch_size, 126, 300)) # float (B, 126, 300) lang_lens = torch.tensor(lang_lens, dtype=torch.int64) # float (B, 1) object_cats = torch.tensor(object_cats) # float (B, ) object_ids = torch.tensor(object_ids, dtype=torch.int64) load_time = torch.from_numpy(np.array(time.time() - start))[None] return { 'locs': locs, 'locs_float': locs_float, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, 'spatial_shape': spatial_shape, 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, 'offsets': batch_offsets, "lang_feat": lang_feats, "lang_len": lang_lens, 'object_id': object_ids, "load_time": load_time, "object_cat": object_cats }