def test_fps():
if not torch.cuda.is_available():
pytest.skip()
xyz = torch.tensor([[[-0.2748, 1.0020, -1.1674], [0.1015, 1.3952, -1.2681],
[-0.8070, 2.4137,
-0.5845], [-1.0001, 2.1982, -0.5859],
[0.3841, 1.8983, -0.7431]],
[[-1.0696, 3.0758,
-0.1899], [-0.2559, 3.5521, -0.1402],
[0.8164, 4.0081, -0.1839], [-1.1000, 3.0213, -0.8205],
[-0.0518, 3.7251, -0.3950]]]).cuda()
idx = furthest_point_sample(xyz, 3)
expected_idx = torch.tensor([[0, 2, 4], [0, 2, 1]]).cuda()
assert torch.all(idx == expected_idx)
def forward(self, feat_dict, sample_mod):
"""Forward pass.
Note:
The forward of VoteHead is devided into 4 steps:
1. Generate vote_points from seed_points.
2. Aggregate vote_points.
3. Predict bbox and score.
4. Decode predictions.
Args:
feat_dict (dict): Feature dict from backbone.
sample_mod (str): Sample mode for vote aggregation layer.
valid modes are "vote", "seed", "random" and "spec".
Returns:
dict: Predictions of vote head.
"""
assert sample_mod in ['vote', 'seed', 'random', 'spec']
seed_points, seed_features, seed_indices = self._extract_input(
feat_dict)
# 1. generate vote_points from seed_points
vote_points, vote_features, vote_offset = self.vote_module(
seed_points, seed_features)
results = dict(
seed_points=seed_points,
seed_indices=seed_indices,
vote_points=vote_points,
vote_features=vote_features,
vote_offset=vote_offset)
# 2. aggregate vote_points
if sample_mod == 'vote':
# use fps in vote_aggregation
aggregation_inputs = dict(
points_xyz=vote_points, features=vote_features)
elif sample_mod == 'seed':
# FPS on seed and choose the votes corresponding to the seeds
sample_indices = furthest_point_sample(seed_points,
self.num_proposal)
aggregation_inputs = dict(
points_xyz=vote_points,
features=vote_features,
indices=sample_indices)
elif sample_mod == 'random':
# Random sampling from the votes
batch_size, num_seed = seed_points.shape[:2]
sample_indices = seed_points.new_tensor(
torch.randint(0, num_seed, (batch_size, self.num_proposal)),
dtype=torch.int32)
aggregation_inputs = dict(
points_xyz=vote_points,
features=vote_features,
indices=sample_indices)
elif sample_mod == 'spec':
# Specify the new center in vote_aggregation
aggregation_inputs = dict(
points_xyz=seed_points,
features=seed_features,
target_xyz=vote_points)
else:
raise NotImplementedError(
f'Sample mode {sample_mod} is not supported!')
vote_aggregation_ret = self.vote_aggregation(**aggregation_inputs)
aggregated_points, features, aggregated_indices = vote_aggregation_ret
# 3. proposal reasoning
features = self.prop_reason(aggregated_points, features)
results['aggregated_points'] = aggregated_points
results['aggregated_features'] = features
results['aggregated_indices'] = aggregated_indices
# 4. predict bbox and score
cls_predictions, reg_predictions = self.conv_pred(features)
# 5. decode predictions
decode_res = self.bbox_coder.split_pred(cls_predictions,
reg_predictions,
aggregated_points)
results.update(decode_res)
return results
def forward(self, feats_dict, sample_mod):
"""Forward pass.
Args:
feats_dict (dict): Feature dict from backbone.
sample_mod (str): Sample mode for vote aggregation layer.
valid modes are "vote", "seed" and "random".
Returns:
dict: Predictions of primitive head.
"""
assert sample_mod in ['vote', 'seed', 'random']
seed_points = feats_dict['fp_xyz_net0'][-1]
seed_features = feats_dict['hd_feature']
results = {}
primitive_flag = self.flag_conv(seed_features)
primitive_flag = self.flag_pred(primitive_flag)
results['pred_flag_' + self.primitive_mode] = primitive_flag
# 1. generate vote_points from seed_points
vote_points, vote_features, _ = self.vote_module(
seed_points, seed_features)
results['vote_' + self.primitive_mode] = vote_points
results['vote_features_' + self.primitive_mode] = vote_features
# 2. aggregate vote_points
if sample_mod == 'vote':
# use fps in vote_aggregation
sample_indices = None
elif sample_mod == 'seed':
# FPS on seed and choose the votes corresponding to the seeds
sample_indices = furthest_point_sample(seed_points,
self.num_proposal)
elif sample_mod == 'random':
# Random sampling from the votes
batch_size, num_seed = seed_points.shape[:2]
sample_indices = torch.randint(
0,
num_seed, (batch_size, self.num_proposal),
dtype=torch.int32,
device=seed_points.device)
else:
raise NotImplementedError('Unsupported sample mod!')
vote_aggregation_ret = self.vote_aggregation(vote_points,
vote_features,
sample_indices)
aggregated_points, features, aggregated_indices = vote_aggregation_ret
results['aggregated_points_' + self.primitive_mode] = aggregated_points
results['aggregated_features_' + self.primitive_mode] = features
results['aggregated_indices_' +
self.primitive_mode] = aggregated_indices
# 3. predict primitive offsets and semantic information
predictions = self.conv_pred(features)
# 4. decode predictions
decode_ret = self.primitive_decode_scores(predictions,
aggregated_points)
results.update(decode_ret)
center, pred_ind = self.get_primitive_center(
primitive_flag, decode_ret['center_' + self.primitive_mode])
results['pred_' + self.primitive_mode + '_ind'] = pred_ind
results['pred_' + self.primitive_mode + '_center'] = center
return results
def forward(self, feat_dict, sample_mod):
"""Forward pass.
Note:
The forward of VoteHead is devided into 4 steps:
1. Generate vote_points from seed_points.
2. Aggregate vote_points.
3. Predict bbox and score.
4. Decode predictions.
Args:
feat_dict (dict): Feature dict from backbone.
sample_mod (str): Sample mode for vote aggregation layer.
valid modes are "vote", "seed" and "random".
Returns:
dict: Predictions of vote head.
"""
assert sample_mod in ['vote', 'seed', 'random']
seed_points = feat_dict['fp_xyz'][-1]
seed_features = feat_dict['fp_features'][-1]
seed_indices = feat_dict['fp_indices'][-1]
# 1. generate vote_points from seed_points
vote_points, vote_features = self.vote_module(seed_points,
seed_features)
results = dict(seed_points=seed_points,
seed_indices=seed_indices,
vote_points=vote_points,
vote_features=vote_features)
# 2. aggregate vote_points
if sample_mod == 'vote':
# use fps in vote_aggregation
sample_indices = None
elif sample_mod == 'seed':
# FPS on seed and choose the votes corresponding to the seeds
sample_indices = furthest_point_sample(seed_points,
self.num_proposal)
elif sample_mod == 'random':
# Random sampling from the votes
batch_size, num_seed = seed_points.shape[:2]
sample_indices = seed_points.new_tensor(torch.randint(
0, num_seed, (batch_size, self.num_proposal)),
dtype=torch.int32)
else:
raise NotImplementedError
vote_aggregation_ret = self.vote_aggregation(vote_points,
vote_features,
sample_indices)
aggregated_points, features, aggregated_indices = vote_aggregation_ret
results['aggregated_points'] = aggregated_points
results['aggregated_indices'] = aggregated_indices
# 3. predict bbox and score
predictions = self.conv_pred(features)
# 4. decode predictions
decode_res = self.bbox_coder.split_pred(predictions, aggregated_points)
results.update(decode_res)
return results
