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
}
