Skip to content

GitLab

Commit 202d767d authored by Shaoshuai Shi's avatar Shaoshuai Shi
Browse files

add config DISABLE_PART for PartA2Head

parent 8ec2a2d4
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
pcdet/models/roi_heads/partA2_head.py
View file @ 202d767d
...@@ -118,7 +118,8 @@ class PartA2FCHead(RoIHeadTemplate): ...@@ -118,7 +118,8 @@ class PartA2FCHead(RoIHeadTemplate):
point_coords = batch_dict['point_coords'][:, 1:4] point_coords = batch_dict['point_coords'][:, 1:4]
point_features = batch_dict['point_features'] point_features = batch_dict['point_features']
part_features = torch.cat(( part_features = torch.cat((
batch_dict['point_part_offset'], batch_dict['point_cls_scores'].view(-1, 1).detach() batch_dict['point_part_offset'] if not self.model_cfg.get('DISABLE_PART', False) else point_coords,
batch_dict['point_cls_scores'].view(-1, 1).detach()
), dim=1) ), dim=1)
part_features[part_features[:, -1] < self.model_cfg.SEG_MASK_SCORE_THRESH, 0:3] = 0 part_features[part_features[:, -1] < self.model_cfg.SEG_MASK_SCORE_THRESH, 0:3] = 0
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment