def __init__(self,
num_classes,
in_channels,
bbox_coder,
num_decoder_layers,
transformerlayers,
decoder_self_posembeds=dict(type='ConvBNPositionalEncoding',
input_channel=6,
num_pos_feats=288),
decoder_cross_posembeds=dict(type='ConvBNPositionalEncoding',
input_channel=3,
num_pos_feats=288),
train_cfg=None,
test_cfg=None,
num_proposal=128,
pred_layer_cfg=None,
size_cls_agnostic=True,
gt_per_seed=3,
sampling_objectness_loss=None,
objectness_loss=None,
center_loss=None,
dir_class_loss=None,
dir_res_loss=None,
size_class_loss=None,
size_res_loss=None,
size_reg_loss=None,
semantic_loss=None,
init_cfg=None):
super(GroupFree3DHead, self).__init__(init_cfg=init_cfg)
self.num_classes = num_classes
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.num_proposal = num_proposal
self.in_channels = in_channels
self.num_decoder_layers = num_decoder_layers
self.size_cls_agnostic = size_cls_agnostic
self.gt_per_seed = gt_per_seed
# Transformer decoder layers
if isinstance(transformerlayers, ConfigDict):
transformerlayers = [
copy.deepcopy(transformerlayers)
for _ in range(num_decoder_layers)
]
else:
assert isinstance(transformerlayers, list) and \
len(transformerlayers) == num_decoder_layers
self.decoder_layers = nn.ModuleList()
for i in range(self.num_decoder_layers):
self.decoder_layers.append(
build_transformer_layer(transformerlayers[i]))
self.embed_dims = self.decoder_layers[0].embed_dims
assert self.embed_dims == decoder_self_posembeds['num_pos_feats']
assert self.embed_dims == decoder_cross_posembeds['num_pos_feats']
# bbox_coder
self.bbox_coder = build_bbox_coder(bbox_coder)
self.num_sizes = self.bbox_coder.num_sizes
self.num_dir_bins = self.bbox_coder.num_dir_bins
# Initial object candidate sampling
self.gsample_module = GeneralSamplingModule()
self.fps_module = Points_Sampler([self.num_proposal])
self.points_obj_cls = PointsObjClsModule(self.in_channels)
self.fp16_enabled = False
# initial candidate prediction
self.conv_pred = BaseConvBboxHead(
**pred_layer_cfg,
num_cls_out_channels=self._get_cls_out_channels(),
num_reg_out_channels=self._get_reg_out_channels())
# query proj and key proj
self.decoder_query_proj = nn.Conv1d(self.embed_dims,
self.embed_dims,
kernel_size=1)
self.decoder_key_proj = nn.Conv1d(self.embed_dims,
self.embed_dims,
kernel_size=1)
# query position embed
self.decoder_self_posembeds = nn.ModuleList()
for _ in range(self.num_decoder_layers):
self.decoder_self_posembeds.append(
build_positional_encoding(decoder_self_posembeds))
# key position embed
self.decoder_cross_posembeds = nn.ModuleList()
for _ in range(self.num_decoder_layers):
self.decoder_cross_posembeds.append(
build_positional_encoding(decoder_cross_posembeds))
# Prediction Head
self.prediction_heads = nn.ModuleList()
for i in range(self.num_decoder_layers):
self.prediction_heads.append(
BaseConvBboxHead(
**pred_layer_cfg,
num_cls_out_channels=self._get_cls_out_channels(),
num_reg_out_channels=self._get_reg_out_channels()))
self.sampling_objectness_loss = build_loss(sampling_objectness_loss)
self.objectness_loss = build_loss(objectness_loss)
self.center_loss = build_loss(center_loss)
self.dir_res_loss = build_loss(dir_res_loss)
self.dir_class_loss = build_loss(dir_class_loss)
self.semantic_loss = build_loss(semantic_loss)
if self.size_cls_agnostic:
self.size_reg_loss = build_loss(size_reg_loss)
else:
self.size_res_loss = build_loss(size_res_loss)
self.size_class_loss = build_loss(size_class_loss)
def __init__(self,
num_point: int,
radii: List[float],
sample_nums: List[int],
mlp_channels: List[List[int]],
fps_mod: List[str] = ['D-FPS'],
fps_sample_range_list: List[int] = [-1],
dilated_group: bool = False,
norm_cfg: dict = dict(type='BN2d'),
use_xyz: bool = True,
pool_mod='max',
normalize_xyz: bool = False,
bias='auto'):
super().__init__()
assert len(radii) == len(sample_nums) == len(mlp_channels)
assert pool_mod in ['max', 'avg']
assert isinstance(fps_mod, list) or isinstance(fps_mod, tuple)
assert isinstance(fps_sample_range_list, list) or isinstance(
fps_sample_range_list, tuple)
assert len(fps_mod) == len(fps_sample_range_list)
if isinstance(mlp_channels, tuple):
mlp_channels = list(map(list, mlp_channels))
if isinstance(num_point, int):
self.num_point = [num_point]
elif isinstance(num_point, list) or isinstance(num_point, tuple):
self.num_point = num_point
else:
raise NotImplementedError('Error type of num_point!')
self.pool_mod = pool_mod
self.groupers = nn.ModuleList()
self.mlps = nn.ModuleList()
self.fps_mod_list = fps_mod
self.fps_sample_range_list = fps_sample_range_list
self.points_sampler = Points_Sampler(self.num_point, self.fps_mod_list,
self.fps_sample_range_list)
for i in range(len(radii)):
radius = radii[i]
sample_num = sample_nums[i]
if num_point is not None:
if dilated_group and i != 0:
min_radius = radii[i - 1]
else:
min_radius = 0
grouper = QueryAndGroup(radius,
sample_num,
min_radius=min_radius,
use_xyz=use_xyz,
normalize_xyz=normalize_xyz)
else:
grouper = GroupAll(use_xyz)
self.groupers.append(grouper)
mlp_spec = mlp_channels[i]
if use_xyz:
mlp_spec[0] += 3
mlp = nn.Sequential()
for i in range(len(mlp_spec) - 1):
mlp.add_module(
f'layer{i}',
ConvModule(mlp_spec[i],
mlp_spec[i + 1],
kernel_size=(1, 1),
stride=(1, 1),
conv_cfg=dict(type='Conv2d'),
act_cfg=dict(type='LeakyReLU',
negative_slope=0.2),
norm_cfg=norm_cfg,
bias=bias))
self.mlps.append(mlp)
def __init__(self,
num_point,
radii,
sample_nums,
mlp_channels,
fps_mod=['D-FPS'],
fps_sample_range_list=[-1],
dilated_group=False,
use_xyz=True,
pool_mod='max',
normalize_xyz=False,
grouper_return_grouped_xyz=False,
grouper_return_grouped_idx=False):
super(BasePointSAModule, self).__init__()
assert len(radii) == len(sample_nums) == len(mlp_channels)
assert pool_mod in ['max', 'avg']
assert isinstance(fps_mod, list) or isinstance(fps_mod, tuple)
assert isinstance(fps_sample_range_list, list) or isinstance(
fps_sample_range_list, tuple)
assert len(fps_mod) == len(fps_sample_range_list)
if isinstance(mlp_channels, tuple):
mlp_channels = list(map(list, mlp_channels))
self.mlp_channels = mlp_channels
if isinstance(num_point, int):
self.num_point = [num_point]
elif isinstance(num_point, list) or isinstance(num_point, tuple):
self.num_point = num_point
else:
raise NotImplementedError('Error type of num_point!')
self.pool_mod = pool_mod
self.groupers = nn.ModuleList()
self.mlps = nn.ModuleList()
self.fps_mod_list = fps_mod
self.fps_sample_range_list = fps_sample_range_list
self.points_sampler = Points_Sampler(self.num_point, self.fps_mod_list,
self.fps_sample_range_list)
for i in range(len(radii)):
radius = radii[i]
sample_num = sample_nums[i]
if num_point is not None:
if dilated_group and i != 0:
min_radius = radii[i - 1]
else:
min_radius = 0
grouper = QueryAndGroup(
radius,
sample_num,
min_radius=min_radius,
use_xyz=use_xyz,
normalize_xyz=normalize_xyz,
return_grouped_xyz=grouper_return_grouped_xyz,
return_grouped_idx=grouper_return_grouped_idx)
else:
grouper = GroupAll(use_xyz)
self.groupers.append(grouper)