885 lines
33 KiB
Python
885 lines
33 KiB
Python
|
import copy
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
from mmcv.cnn import ConvModule, build_conv_layer
|
|||
|
|
from mmcv.runner import BaseModule, force_fp32
|
|||
|
|
from torch import nn
|
|||
|
|
|
|||
|
|
from mmdet3d.core import circle_nms, draw_heatmap_gaussian, gaussian_radius, xywhr2xyxyr
|
|||
|
|
from mmdet3d.models import builder
|
|||
|
|
from mmdet3d.models.builder import HEADS, build_loss
|
|||
|
|
from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu
|
|||
|
|
from mmdet.core import build_bbox_coder, multi_apply
|
|||
|
|
|
|||
|
|
|
|||
|
|
def clip_sigmoid(x: torch.Tensor, eps: float = 1e-4) -> torch.Tensor:
|
|||
|
|
return torch.clamp(x.sigmoid_(), min=eps, max=1 - eps)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@HEADS.register_module()
|
|||
|
|
class SeparateHead(BaseModule):
|
|||
|
|
"""SeparateHead for CenterHead.
|
|||
|
|
Args:
|
|||
|
|
in_channels (int): Input channels for conv_layer.
|
|||
|
|
heads (dict): Conv information.
|
|||
|
|
head_conv (int): Output channels.
|
|||
|
|
Default: 64.
|
|||
|
|
final_kernal (int): Kernal size for the last conv layer.
|
|||
|
|
Deafult: 1.
|
|||
|
|
init_bias (float): Initial bias. Default: -2.19.
|
|||
|
|
conv_cfg (dict): Config of conv layer.
|
|||
|
|
Default: dict(type='Conv2d')
|
|||
|
|
norm_cfg (dict): Config of norm layer.
|
|||
|
|
Default: dict(type='BN2d').
|
|||
|
|
bias (str): Type of bias. Default: 'auto'.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
in_channels,
|
|||
|
|
heads,
|
|||
|
|
head_conv=64,
|
|||
|
|
final_kernel=1,
|
|||
|
|
init_bias=-2.19,
|
|||
|
|
conv_cfg=dict(type="Conv2d"),
|
|||
|
|
norm_cfg=dict(type="BN2d"),
|
|||
|
|
bias="auto",
|
|||
|
|
init_cfg=None,
|
|||
|
|
**kwargs,
|
|||
|
|
):
|
|||
|
|
assert (
|
|||
|
|
init_cfg is None
|
|||
|
|
), "To prevent abnormal initialization behavior, init_cfg is not allowed to be set"
|
|||
|
|
super(SeparateHead, self).__init__(init_cfg=init_cfg)
|
|||
|
|
self.heads = heads
|
|||
|
|
self.init_bias = init_bias
|
|||
|
|
for head in self.heads:
|
|||
|
|
classes, num_conv = self.heads[head]
|
|||
|
|
|
|||
|
|
conv_layers = []
|
|||
|
|
c_in = in_channels
|
|||
|
|
for i in range(num_conv - 1):
|
|||
|
|
conv_layers.append(
|
|||
|
|
ConvModule(
|
|||
|
|
c_in,
|
|||
|
|
head_conv,
|
|||
|
|
kernel_size=final_kernel,
|
|||
|
|
stride=1,
|
|||
|
|
padding=final_kernel // 2,
|
|||
|
|
bias=bias,
|
|||
|
|
conv_cfg=conv_cfg,
|
|||
|
|
norm_cfg=norm_cfg,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
c_in = head_conv
|
|||
|
|
|
|||
|
|
conv_layers.append(
|
|||
|
|
build_conv_layer(
|
|||
|
|
conv_cfg,
|
|||
|
|
head_conv,
|
|||
|
|
classes,
|
|||
|
|
kernel_size=final_kernel,
|
|||
|
|
stride=1,
|
|||
|
|
padding=final_kernel // 2,
|
|||
|
|
bias=True,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
conv_layers = nn.Sequential(*conv_layers)
|
|||
|
|
|
|||
|
|
self.__setattr__(head, conv_layers)
|
|||
|
|
|
|||
|
|
if init_cfg is None:
|
|||
|
|
self.init_cfg = dict(type="Kaiming", layer="Conv2d")
|
|||
|
|
|
|||
|
|
def init_weights(self):
|
|||
|
|
"""Initialize weights."""
|
|||
|
|
super().init_weights()
|
|||
|
|
for head in self.heads:
|
|||
|
|
if head == "heatmap":
|
|||
|
|
self.__getattr__(head)[-1].bias.data.fill_(self.init_bias)
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
"""Forward function for SepHead.
|
|||
|
|
Args:
|
|||
|
|
x (torch.Tensor): Input feature map with the shape of
|
|||
|
|
[B, 512, 128, 128].
|
|||
|
|
Returns:
|
|||
|
|
dict[str: torch.Tensor]: contains the following keys:
|
|||
|
|
-reg （torch.Tensor): 2D regression value with the \
|
|||
|
|
shape of [B, 2, H, W].
|
|||
|
|
-height (torch.Tensor): Height value with the \
|
|||
|
|
shape of [B, 1, H, W].
|
|||
|
|
-dim (torch.Tensor): Size value with the shape \
|
|||
|
|
of [B, 3, H, W].
|
|||
|
|
-rot (torch.Tensor): Rotation value with the \
|
|||
|
|
shape of [B, 2, H, W].
|
|||
|
|
-vel (torch.Tensor): Velocity value with the \
|
|||
|
|
shape of [B, 2, H, W].
|
|||
|
|
-heatmap (torch.Tensor): Heatmap with the shape of \
|
|||
|
|
[B, N, H, W].
|
|||
|
|
"""
|
|||
|
|
ret_dict = dict()
|
|||
|
|
for head in self.heads:
|
|||
|
|
ret_dict[head] = self.__getattr__(head)(x)
|
|||
|
|
|
|||
|
|
return ret_dict
|
|||
|
|
|
|||
|
|
|
|||
|
|
@HEADS.register_module()
|
|||
|
|
class DCNSeparateHead(BaseModule):
|
|||
|
|
r"""DCNSeparateHead for CenterHead.
|
|||
|
|
.. code-block:: none
|
|||
|
|
/-----> DCN for heatmap task -----> heatmap task.
|
|||
|
|
feature
|
|||
|
|
\-----> DCN for regression tasks -----> regression tasks
|
|||
|
|
Args:
|
|||
|
|
in_channels (int): Input channels for conv_layer.
|
|||
|
|
heads (dict): Conv information.
|
|||
|
|
dcn_config (dict): Config of dcn layer.
|
|||
|
|
num_cls (int): Output channels.
|
|||
|
|
Default: 64.
|
|||
|
|
final_kernal (int): Kernal size for the last conv layer.
|
|||
|
|
Deafult: 1.
|
|||
|
|
init_bias (float): Initial bias. Default: -2.19.
|
|||
|
|
conv_cfg (dict): Config of conv layer.
|
|||
|
|
Default: dict(type='Conv2d')
|
|||
|
|
norm_cfg (dict): Config of norm layer.
|
|||
|
|
Default: dict(type='BN2d').
|
|||
|
|
bias (str): Type of bias. Default: 'auto'.
|
|||
|
|
""" # noqa: W605
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
in_channels,
|
|||
|
|
num_cls,
|
|||
|
|
heads,
|
|||
|
|
dcn_config,
|
|||
|
|
head_conv=64,
|
|||
|
|
final_kernel=1,
|
|||
|
|
init_bias=-2.19,
|
|||
|
|
conv_cfg=dict(type="Conv2d"),
|
|||
|
|
norm_cfg=dict(type="BN2d"),
|
|||
|
|
bias="auto",
|
|||
|
|
init_cfg=None,
|
|||
|
|
**kwargs,
|
|||
|
|
):
|
|||
|
|
assert init_cfg is None, (
|
|||
|
|
"To prevent abnormal initialization "
|
|||
|
|
"behavior, init_cfg is not allowed to be set"
|
|||
|
|
)
|
|||
|
|
super(DCNSeparateHead, self).__init__(init_cfg=init_cfg)
|
|||
|
|
if "heatmap" in heads:
|
|||
|
|
heads.pop("heatmap")
|
|||
|
|
# feature adaptation with dcn
|
|||
|
|
# use separate features for classification / regression
|
|||
|
|
self.feature_adapt_cls = build_conv_layer(dcn_config)
|
|||
|
|
|
|||
|
|
self.feature_adapt_reg = build_conv_layer(dcn_config)
|
|||
|
|
|
|||
|
|
# heatmap prediction head
|
|||
|
|
cls_head = [
|
|||
|
|
ConvModule(
|
|||
|
|
in_channels,
|
|||
|
|
head_conv,
|
|||
|
|
kernel_size=3,
|
|||
|
|
padding=1,
|
|||
|
|
conv_cfg=conv_cfg,
|
|||
|
|
bias=bias,
|
|||
|
|
norm_cfg=norm_cfg,
|
|||
|
|
),
|
|||
|
|
build_conv_layer(
|
|||
|
|
conv_cfg,
|
|||
|
|
head_conv,
|
|||
|
|
num_cls,
|
|||
|
|
kernel_size=3,
|
|||
|
|
stride=1,
|
|||
|
|
padding=1,
|
|||
|
|
bias=bias,
|
|||
|
|
),
|
|||
|
|
]
|
|||
|
|
self.cls_head = nn.Sequential(*cls_head)
|
|||
|
|
self.init_bias = init_bias
|
|||
|
|
# other regression target
|
|||
|
|
self.task_head = SeparateHead(
|
|||
|
|
in_channels,
|
|||
|
|
heads,
|
|||
|
|
head_conv=head_conv,
|
|||
|
|
final_kernel=final_kernel,
|
|||
|
|
bias=bias,
|
|||
|
|
)
|
|||
|
|
if init_cfg is None:
|
|||
|
|
self.init_cfg = dict(type="Kaiming", layer="Conv2d")
|
|||
|
|
|
|||
|
|
def init_weights(self):
|
|||
|
|
"""Initialize weights."""
|
|||
|
|
super().init_weights()
|
|||
|
|
self.cls_head[-1].bias.data.fill_(self.init_bias)
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
"""Forward function for DCNSepHead.
|
|||
|
|
Args:
|
|||
|
|
x (torch.Tensor): Input feature map with the shape of
|
|||
|
|
[B, 512, 128, 128].
|
|||
|
|
Returns:
|
|||
|
|
dict[str: torch.Tensor]: contains the following keys:
|
|||
|
|
-reg （torch.Tensor): 2D regression value with the \
|
|||
|
|
shape of [B, 2, H, W].
|
|||
|
|
-height (torch.Tensor): Height value with the \
|
|||
|
|
shape of [B, 1, H, W].
|
|||
|
|
-dim (torch.Tensor): Size value with the shape \
|
|||
|
|
of [B, 3, H, W].
|
|||
|
|
-rot (torch.Tensor): Rotation value with the \
|
|||
|
|
shape of [B, 2, H, W].
|
|||
|
|
-vel (torch.Tensor): Velocity value with the \
|
|||
|
|
shape of [B, 2, H, W].
|
|||
|
|
-heatmap (torch.Tensor): Heatmap with the shape of \
|
|||
|
|
[B, N, H, W].
|
|||
|
|
"""
|
|||
|
|
center_feat = self.feature_adapt_cls(x)
|
|||
|
|
reg_feat = self.feature_adapt_reg(x)
|
|||
|
|
|
|||
|
|
cls_score = self.cls_head(center_feat)
|
|||
|
|
ret = self.task_head(reg_feat)
|
|||
|
|
ret["heatmap"] = cls_score
|
|||
|
|
|
|||
|
|
return ret
|
|||
|
|
|
|||
|
|
|
|||
|
|
@HEADS.register_module()
|
|||
|
|
class CenterHead(BaseModule):
|
|||
|
|
"""CenterHead for CenterPoint.
|
|||
|
|
Args:
|
|||
|
|
mode (str): Mode of the head. Default: '3d'.
|
|||
|
|
in_channels (list[int] | int): Channels of the input feature map.
|
|||
|
|
Default: [128].
|
|||
|
|
tasks (list[dict]): Task information including class number
|
|||
|
|
and class names. Default: None.
|
|||
|
|
dataset (str): Name of the dataset. Default: 'nuscenes'.
|
|||
|
|
weight (float): Weight for location loss. Default: 0.25.
|
|||
|
|
code_weights (list[int]): Code weights for location loss. Default: [].
|
|||
|
|
common_heads (dict): Conv information for common heads.
|
|||
|
|
Default: dict().
|
|||
|
|
loss_cls (dict): Config of classification loss function.
|
|||
|
|
Default: dict(type='GaussianFocalLoss', reduction='mean').
|
|||
|
|
loss_bbox (dict): Config of regression loss function.
|
|||
|
|
Default: dict(type='L1Loss', reduction='none').
|
|||
|
|
separate_head (dict): Config of separate head. Default: dict(
|
|||
|
|
type='SeparateHead', init_bias=-2.19, final_kernel=3)
|
|||
|
|
share_conv_channel (int): Output channels for share_conv_layer.
|
|||
|
|
Default: 64.
|
|||
|
|
num_heatmap_convs (int): Number of conv layers for heatmap conv layer.
|
|||
|
|
Default: 2.
|
|||
|
|
conv_cfg (dict): Config of conv layer.
|
|||
|
|
Default: dict(type='Conv2d')
|
|||
|
|
norm_cfg (dict): Config of norm layer.
|
|||
|
|
Default: dict(type='BN2d').
|
|||
|
|
bias (str): Type of bias. Default: 'auto'.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
in_channels=[128],
|
|||
|
|
tasks=None,
|
|||
|
|
train_cfg=None,
|
|||
|
|
test_cfg=None,
|
|||
|
|
bbox_coder=None,
|
|||
|
|
common_heads=dict(),
|
|||
|
|
loss_cls=dict(type="GaussianFocalLoss", reduction="mean"),
|
|||
|
|
loss_bbox=dict(type="L1Loss", reduction="none", loss_weight=0.25),
|
|||
|
|
separate_head=dict(type="SeparateHead", init_bias=-2.19, final_kernel=3),
|
|||
|
|
share_conv_channel=64,
|
|||
|
|
num_heatmap_convs=2,
|
|||
|
|
conv_cfg=dict(type="Conv2d"),
|
|||
|
|
norm_cfg=dict(type="BN2d"),
|
|||
|
|
bias="auto",
|
|||
|
|
norm_bbox=True,
|
|||
|
|
init_cfg=None,
|
|||
|
|
):
|
|||
|
|
assert init_cfg is None, (
|
|||
|
|
"To prevent abnormal initialization "
|
|||
|
|
"behavior, init_cfg is not allowed to be set"
|
|||
|
|
)
|
|||
|
|
super(CenterHead, self).__init__(init_cfg=init_cfg)
|
|||
|
|
|
|||
|
|
num_classes = [len(t) for t in tasks]
|
|||
|
|
self.class_names = [t for t in tasks]
|
|||
|
|
self.train_cfg = train_cfg
|
|||
|
|
self.test_cfg = test_cfg
|
|||
|
|
self.in_channels = in_channels
|
|||
|
|
self.num_classes = num_classes
|
|||
|
|
self.norm_bbox = norm_bbox
|
|||
|
|
|
|||
|
|
self.loss_cls = build_loss(loss_cls)
|
|||
|
|
self.loss_bbox = build_loss(loss_bbox)
|
|||
|
|
self.bbox_coder = build_bbox_coder(bbox_coder)
|
|||
|
|
self.num_anchor_per_locs = [n for n in num_classes]
|
|||
|
|
self.fp16_enabled = False
|
|||
|
|
|
|||
|
|
# a shared convolution
|
|||
|
|
self.shared_conv = ConvModule(
|
|||
|
|
in_channels,
|
|||
|
|
share_conv_channel,
|
|||
|
|
kernel_size=3,
|
|||
|
|
padding=1,
|
|||
|
|
conv_cfg=conv_cfg,
|
|||
|
|
norm_cfg=norm_cfg,
|
|||
|
|
bias=bias,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self.task_heads = nn.ModuleList()
|
|||
|
|
|
|||
|
|
for num_cls in num_classes:
|
|||
|
|
heads = copy.deepcopy(common_heads)
|
|||
|
|
heads.update(dict(heatmap=(num_cls, num_heatmap_convs)))
|
|||
|
|
separate_head.update(
|
|||
|
|
in_channels=share_conv_channel, heads=heads, num_cls=num_cls
|
|||
|
|
)
|
|||
|
|
self.task_heads.append(builder.build_head(separate_head))
|
|||
|
|
|
|||
|
|
def forward_single(self, x):
|
|||
|
|
"""Forward function for CenterPoint.
|
|||
|
|
Args:
|
|||
|
|
x (torch.Tensor): Input feature map with the shape of
|
|||
|
|
[B, 512, 128, 128].
|
|||
|
|
Returns:
|
|||
|
|
list[dict]: Output results for tasks.
|
|||
|
|
"""
|
|||
|
|
ret_dicts = []
|
|||
|
|
|
|||
|
|
x = self.shared_conv(x)
|
|||
|
|
|
|||
|
|
for task in self.task_heads:
|
|||
|
|
ret_dicts.append(task(x))
|
|||
|
|
|
|||
|
|
return ret_dicts
|
|||
|
|
|
|||
|
|
def forward(self, feats, metas):
|
|||
|
|
"""Forward pass.
|
|||
|
|
Args:
|
|||
|
|
feats (list[torch.Tensor]): Multi-level features, e.g.,
|
|||
|
|
features produced by FPN.
|
|||
|
|
Returns:
|
|||
|
|
tuple(list[dict]): Output results for tasks.
|
|||
|
|
"""
|
|||
|
|
if isinstance(feats, torch.Tensor):
|
|||
|
|
feats = [feats]
|
|||
|
|
return multi_apply(self.forward_single, feats)
|
|||
|
|
|
|||
|
|
def _gather_feat(self, feat, ind, mask=None):
|
|||
|
|
"""Gather feature map.
|
|||
|
|
Given feature map and index, return indexed feature map.
|
|||
|
|
Args:
|
|||
|
|
feat (torch.tensor): Feature map with the shape of [B, H*W, 10].
|
|||
|
|
ind (torch.Tensor): Index of the ground truth boxes with the
|
|||
|
|
shape of [B, max_obj].
|
|||
|
|
mask (torch.Tensor): Mask of the feature map with the shape
|
|||
|
|
of [B, max_obj]. Default: None.
|
|||
|
|
Returns:
|
|||
|
|
torch.Tensor: Feature map after gathering with the shape
|
|||
|
|
of [B, max_obj, 10].
|
|||
|
|
"""
|
|||
|
|
dim = feat.size(2)
|
|||
|
|
ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
|
|||
|
|
feat = feat.gather(1, ind)
|
|||
|
|
if mask is not None:
|
|||
|
|
mask = mask.unsqueeze(2).expand_as(feat)
|
|||
|
|
feat = feat[mask]
|
|||
|
|
feat = feat.view(-1, dim)
|
|||
|
|
return feat
|
|||
|
|
|
|||
|
|
def get_targets(self, gt_bboxes_3d, gt_labels_3d):
|
|||
|
|
"""Generate targets.
|
|||
|
|
How each output is transformed:
|
|||
|
|
Each nested list is transposed so that all same-index elements in
|
|||
|
|
each sub-list (1, ..., N) become the new sub-lists.
|
|||
|
|
[ [a0, a1, a2, ... ], [b0, b1, b2, ... ], ... ]
|
|||
|
|
==> [ [a0, b0, ... ], [a1, b1, ... ], [a2, b2, ... ] ]
|
|||
|
|
The new transposed nested list is converted into a list of N
|
|||
|
|
tensors generated by concatenating tensors in the new sub-lists.
|
|||
|
|
[ tensor0, tensor1, tensor2, ... ]
|
|||
|
|
Args:
|
|||
|
|
gt_bboxes_3d (list[:obj:`LiDARInstance3DBoxes`]): Ground
|
|||
|
|
truth gt boxes.
|
|||
|
|
gt_labels_3d (list[torch.Tensor]): Labels of boxes.
|
|||
|
|
Returns:
|
|||
|
|
Returns:
|
|||
|
|
tuple[list[torch.Tensor]]: Tuple of target including \
|
|||
|
|
the following results in order.
|
|||
|
|
- list[torch.Tensor]: Heatmap scores.
|
|||
|
|
- list[torch.Tensor]: Ground truth boxes.
|
|||
|
|
- list[torch.Tensor]: Indexes indicating the \
|
|||
|
|
position of the valid boxes.
|
|||
|
|
- list[torch.Tensor]: Masks indicating which \
|
|||
|
|
boxes are valid.
|
|||
|
|
"""
|
|||
|
|
heatmaps, anno_boxes, inds, masks = multi_apply(
|
|||
|
|
self.get_targets_single, gt_bboxes_3d, gt_labels_3d
|
|||
|
|
)
|
|||
|
|
# Transpose heatmaps
|
|||
|
|
heatmaps = list(map(list, zip(*heatmaps)))
|
|||
|
|
heatmaps = [torch.stack(hms_) for hms_ in heatmaps]
|
|||
|
|
# Transpose anno_boxes
|
|||
|
|
anno_boxes = list(map(list, zip(*anno_boxes)))
|
|||
|
|
anno_boxes = [torch.stack(anno_boxes_) for anno_boxes_ in anno_boxes]
|
|||
|
|
# Transpose inds
|
|||
|
|
inds = list(map(list, zip(*inds)))
|
|||
|
|
inds = [torch.stack(inds_) for inds_ in inds]
|
|||
|
|
# Transpose inds
|
|||
|
|
masks = list(map(list, zip(*masks)))
|
|||
|
|
masks = [torch.stack(masks_) for masks_ in masks]
|
|||
|
|
return heatmaps, anno_boxes, inds, masks
|
|||
|
|
|
|||
|
|
def get_targets_single(self, gt_bboxes_3d, gt_labels_3d):
|
|||
|
|
"""Generate training targets for a single sample.
|
|||
|
|
Args:
|
|||
|
|
gt_bboxes_3d (:obj:`LiDARInstance3DBoxes`): Ground truth gt boxes.
|
|||
|
|
gt_labels_3d (torch.Tensor): Labels of boxes.
|
|||
|
|
Returns:
|
|||
|
|
tuple[list[torch.Tensor]]: Tuple of target including \
|
|||
|
|
the following results in order.
|
|||
|
|
- list[torch.Tensor]: Heatmap scores.
|
|||
|
|
- list[torch.Tensor]: Ground truth boxes.
|
|||
|
|
- list[torch.Tensor]: Indexes indicating the position \
|
|||
|
|
of the valid boxes.
|
|||
|
|
- list[torch.Tensor]: Masks indicating which boxes \
|
|||
|
|
are valid.
|
|||
|
|
"""
|
|||
|
|
device = gt_labels_3d.device
|
|||
|
|
gt_bboxes_3d = torch.cat(
|
|||
|
|
(gt_bboxes_3d.gravity_center, gt_bboxes_3d.tensor[:, 3:]), dim=1
|
|||
|
|
).to(device)
|
|||
|
|
max_objs = self.train_cfg["max_objs"] * self.train_cfg["dense_reg"]
|
|||
|
|
grid_size = torch.tensor(self.train_cfg["grid_size"])
|
|||
|
|
pc_range = torch.tensor(self.train_cfg["point_cloud_range"])
|
|||
|
|
voxel_size = torch.tensor(self.train_cfg["voxel_size"])
|
|||
|
|
|
|||
|
|
feature_map_size = torch.div(
|
|||
|
|
grid_size[:2],
|
|||
|
|
self.train_cfg["out_size_factor"],
|
|||
|
|
rounding_mode="trunc",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# reorganize the gt_dict by tasks
|
|||
|
|
task_masks = []
|
|||
|
|
flag = 0
|
|||
|
|
for class_name in self.class_names:
|
|||
|
|
task_masks.append(
|
|||
|
|
[
|
|||
|
|
torch.where(gt_labels_3d == class_name.index(i) + flag)
|
|||
|
|
for i in class_name
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
flag += len(class_name)
|
|||
|
|
|
|||
|
|
task_boxes = []
|
|||
|
|
task_classes = []
|
|||
|
|
flag2 = 0
|
|||
|
|
for idx, mask in enumerate(task_masks):
|
|||
|
|
task_box = []
|
|||
|
|
task_class = []
|
|||
|
|
for m in mask:
|
|||
|
|
task_box.append(gt_bboxes_3d[m])
|
|||
|
|
# 0 is background for each task, so we need to add 1 here.
|
|||
|
|
task_class.append(gt_labels_3d[m] + 1 - flag2)
|
|||
|
|
task_boxes.append(torch.cat(task_box, axis=0).to(device))
|
|||
|
|
task_classes.append(torch.cat(task_class).long().to(device))
|
|||
|
|
flag2 += len(mask)
|
|||
|
|
draw_gaussian = draw_heatmap_gaussian
|
|||
|
|
heatmaps, anno_boxes, inds, masks = [], [], [], []
|
|||
|
|
|
|||
|
|
for idx, task_head in enumerate(self.task_heads):
|
|||
|
|
heatmap = gt_bboxes_3d.new_zeros(
|
|||
|
|
(len(self.class_names[idx]), feature_map_size[1], feature_map_size[0])
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
anno_box = gt_bboxes_3d.new_zeros((max_objs, 10), dtype=torch.float32)
|
|||
|
|
|
|||
|
|
ind = gt_labels_3d.new_zeros((max_objs), dtype=torch.int64)
|
|||
|
|
mask = gt_bboxes_3d.new_zeros((max_objs), dtype=torch.uint8)
|
|||
|
|
|
|||
|
|
num_objs = min(task_boxes[idx].shape[0], max_objs)
|
|||
|
|
|
|||
|
|
for k in range(num_objs):
|
|||
|
|
cls_id = task_classes[idx][k] - 1
|
|||
|
|
|
|||
|
|
width = task_boxes[idx][k][3]
|
|||
|
|
length = task_boxes[idx][k][4]
|
|||
|
|
width = width / voxel_size[0] / self.train_cfg["out_size_factor"]
|
|||
|
|
length = length / voxel_size[1] / self.train_cfg["out_size_factor"]
|
|||
|
|
|
|||
|
|
if width > 0 and length > 0:
|
|||
|
|
radius = gaussian_radius(
|
|||
|
|
(length, width), min_overlap=self.train_cfg["gaussian_overlap"]
|
|||
|
|
)
|
|||
|
|
radius = max(self.train_cfg["min_radius"], int(radius))
|
|||
|
|
|
|||
|
|
# be really careful for the coordinate system of
|
|||
|
|
# your box annotation.
|
|||
|
|
x, y, z = (
|
|||
|
|
task_boxes[idx][k][0],
|
|||
|
|
task_boxes[idx][k][1],
|
|||
|
|
task_boxes[idx][k][2],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
coor_x = (
|
|||
|
|
(x - pc_range[0])
|
|||
|
|
/ voxel_size[0]
|
|||
|
|
/ self.train_cfg["out_size_factor"]
|
|||
|
|
)
|
|||
|
|
coor_y = (
|
|||
|
|
(y - pc_range[1])
|
|||
|
|
/ voxel_size[1]
|
|||
|
|
/ self.train_cfg["out_size_factor"]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
center = torch.tensor(
|
|||
|
|
[coor_x, coor_y], dtype=torch.float32, device=device
|
|||
|
|
)
|
|||
|
|
center_int = center.to(torch.int32)
|
|||
|
|
|
|||
|
|
# throw out not in range objects to avoid out of array
|
|||
|
|
# area when creating the heatmap
|
|||
|
|
if not (
|
|||
|
|
0 <= center_int[0] < feature_map_size[0]
|
|||
|
|
and 0 <= center_int[1] < feature_map_size[1]
|
|||
|
|
):
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
draw_gaussian(heatmap[cls_id], center_int[[1, 0]], radius)
|
|||
|
|
new_idx = k
|
|||
|
|
x, y = center_int[0], center_int[1]
|
|||
|
|
|
|||
|
|
assert (
|
|||
|
|
x * feature_map_size[1] + y
|
|||
|
|
< feature_map_size[0] * feature_map_size[1]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
ind[new_idx] = x * feature_map_size[1] + y
|
|||
|
|
|
|||
|
|
mask[new_idx] = 1
|
|||
|
|
# TODO: support other outdoor dataset
|
|||
|
|
vx, vy = task_boxes[idx][k][7:]
|
|||
|
|
rot = task_boxes[idx][k][6]
|
|||
|
|
box_dim = task_boxes[idx][k][3:6]
|
|||
|
|
if self.norm_bbox:
|
|||
|
|
box_dim = box_dim.log()
|
|||
|
|
anno_box[new_idx] = torch.cat(
|
|||
|
|
[
|
|||
|
|
center - torch.tensor([x, y], device=device),
|
|||
|
|
z.unsqueeze(0),
|
|||
|
|
box_dim,
|
|||
|
|
torch.sin(rot).unsqueeze(0),
|
|||
|
|
torch.cos(rot).unsqueeze(0),
|
|||
|
|
vx.unsqueeze(0),
|
|||
|
|
vy.unsqueeze(0),
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
heatmaps.append(heatmap)
|
|||
|
|
anno_boxes.append(anno_box)
|
|||
|
|
masks.append(mask)
|
|||
|
|
inds.append(ind)
|
|||
|
|
return heatmaps, anno_boxes, inds, masks
|
|||
|
|
|
|||
|
|
@force_fp32(apply_to=("preds_dicts"))
|
|||
|
|
def loss(self, gt_bboxes_3d, gt_labels_3d, preds_dicts, **kwargs):
|
|||
|
|
"""Loss function for CenterHead.
|
|||
|
|
Args:
|
|||
|
|
gt_bboxes_3d (list[:obj:`LiDARInstance3DBoxes`]): Ground
|
|||
|
|
truth gt boxes.
|
|||
|
|
gt_labels_3d (list[torch.Tensor]): Labels of boxes.
|
|||
|
|
preds_dicts (dict): Output of forward function.
|
|||
|
|
Returns:
|
|||
|
|
dict[str:torch.Tensor]: Loss of heatmap and bbox of each task.
|
|||
|
|
"""
|
|||
|
|
heatmaps, anno_boxes, inds, masks = self.get_targets(gt_bboxes_3d, gt_labels_3d)
|
|||
|
|
loss_dict = dict()
|
|||
|
|
for task_id, preds_dict in enumerate(preds_dicts):
|
|||
|
|
# heatmap focal loss
|
|||
|
|
preds_dict[0]["heatmap"] = clip_sigmoid(preds_dict[0]["heatmap"])
|
|||
|
|
num_pos = heatmaps[task_id].eq(1).float().sum().item()
|
|||
|
|
loss_heatmap = self.loss_cls(
|
|||
|
|
preds_dict[0]["heatmap"], heatmaps[task_id], avg_factor=max(num_pos, 1)
|
|||
|
|
)
|
|||
|
|
target_box = anno_boxes[task_id]
|
|||
|
|
# reconstruct the anno_box from multiple reg heads
|
|||
|
|
preds_dict[0]["anno_box"] = torch.cat(
|
|||
|
|
(
|
|||
|
|
preds_dict[0]["reg"],
|
|||
|
|
preds_dict[0]["height"],
|
|||
|
|
preds_dict[0]["dim"],
|
|||
|
|
preds_dict[0]["rot"],
|
|||
|
|
preds_dict[0]["vel"],
|
|||
|
|
),
|
|||
|
|
dim=1,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Regression loss for dimension, offset, height, rotation
|
|||
|
|
ind = inds[task_id]
|
|||
|
|
num = masks[task_id].float().sum()
|
|||
|
|
pred = preds_dict[0]["anno_box"].permute(0, 2, 3, 1).contiguous()
|
|||
|
|
pred = pred.view(pred.size(0), -1, pred.size(3))
|
|||
|
|
pred = self._gather_feat(pred, ind)
|
|||
|
|
mask = masks[task_id].unsqueeze(2).expand_as(target_box).float()
|
|||
|
|
isnotnan = (~torch.isnan(target_box)).float()
|
|||
|
|
mask *= isnotnan
|
|||
|
|
|
|||
|
|
code_weights = self.train_cfg.get("code_weights", None)
|
|||
|
|
bbox_weights = mask * mask.new_tensor(code_weights)
|
|||
|
|
loss_bbox = self.loss_bbox(
|
|||
|
|
pred, target_box, bbox_weights, avg_factor=(num + 1e-4)
|
|||
|
|
)
|
|||
|
|
loss_dict[f"heatmap/task{task_id}"] = loss_heatmap
|
|||
|
|
loss_dict[f"bbox/task{task_id}"] = loss_bbox
|
|||
|
|
return loss_dict
|
|||
|
|
|
|||
|
|
@force_fp32(apply_to=("preds_dicts"))
|
|||
|
|
def get_bboxes(self, preds_dicts, metas, img=None, rescale=False):
|
|||
|
|
"""Generate bboxes from bbox head predictions.
|
|||
|
|
Args:
|
|||
|
|
preds_dicts (tuple[list[dict]]): Prediction results.
|
|||
|
|
metas (list[dict]): Point cloud and image's meta info.
|
|||
|
|
Returns:
|
|||
|
|
list[dict]: Decoded bbox, scores and labels after nms.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
if not isinstance(self.test_cfg["nms_type"], list):
|
|||
|
|
nms_types = [self.test_cfg["nms_type"] for _ in range(len(preds_dicts))]
|
|||
|
|
else:
|
|||
|
|
nms_types = self.test_cfg["nms_type"]
|
|||
|
|
|
|||
|
|
if "nms_scale" in self.test_cfg:
|
|||
|
|
if not isinstance(self.test_cfg["nms_scale"], list):
|
|||
|
|
nms_scales = [
|
|||
|
|
[
|
|||
|
|
self.test_cfg["nms_scale"]
|
|||
|
|
for _ in range(self.num_classes[task_id])
|
|||
|
|
]
|
|||
|
|
for task_id in range(len(preds_dicts))
|
|||
|
|
]
|
|||
|
|
else:
|
|||
|
|
nms_scales = self.test_cfg["nms_scale"]
|
|||
|
|
else:
|
|||
|
|
nms_scales = [
|
|||
|
|
[1.0 for _ in range(self.num_classes[task_id])]
|
|||
|
|
for task_id in range(len(preds_dicts))
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
rets = []
|
|||
|
|
for task_id, preds_dict in enumerate(preds_dicts):
|
|||
|
|
num_class_with_bg = self.num_classes[task_id]
|
|||
|
|
batch_size = preds_dict[0]["heatmap"].shape[0]
|
|||
|
|
batch_heatmap = preds_dict[0]["heatmap"].sigmoid()
|
|||
|
|
|
|||
|
|
batch_reg = preds_dict[0]["reg"]
|
|||
|
|
batch_hei = preds_dict[0]["height"]
|
|||
|
|
|
|||
|
|
if self.norm_bbox:
|
|||
|
|
batch_dim = torch.exp(preds_dict[0]["dim"])
|
|||
|
|
else:
|
|||
|
|
batch_dim = preds_dict[0]["dim"]
|
|||
|
|
|
|||
|
|
batch_rots = preds_dict[0]["rot"][:, 0].unsqueeze(1)
|
|||
|
|
batch_rotc = preds_dict[0]["rot"][:, 1].unsqueeze(1)
|
|||
|
|
|
|||
|
|
if "vel" in preds_dict[0]:
|
|||
|
|
batch_vel = preds_dict[0]["vel"]
|
|||
|
|
else:
|
|||
|
|
batch_vel = None
|
|||
|
|
temp = self.bbox_coder.decode(
|
|||
|
|
batch_heatmap,
|
|||
|
|
batch_rots,
|
|||
|
|
batch_rotc,
|
|||
|
|
batch_hei,
|
|||
|
|
batch_dim,
|
|||
|
|
batch_vel,
|
|||
|
|
reg=batch_reg,
|
|||
|
|
task_id=task_id,
|
|||
|
|
)
|
|||
|
|
batch_reg_preds = [box["bboxes"] for box in temp]
|
|||
|
|
batch_cls_preds = [box["scores"] for box in temp]
|
|||
|
|
batch_cls_labels = [box["labels"] for box in temp]
|
|||
|
|
if nms_types[task_id] == "circle":
|
|||
|
|
ret_task = []
|
|||
|
|
for i in range(batch_size):
|
|||
|
|
boxes3d = temp[i]["bboxes"]
|
|||
|
|
scores = temp[i]["scores"]
|
|||
|
|
labels = temp[i]["labels"]
|
|||
|
|
centers = boxes3d[:, [0, 1]]
|
|||
|
|
boxes = torch.cat([centers, scores.view(-1, 1)], dim=1)
|
|||
|
|
keep = torch.tensor(
|
|||
|
|
circle_nms(
|
|||
|
|
boxes.detach().cpu().numpy(),
|
|||
|
|
self.test_cfg["min_radius"][task_id],
|
|||
|
|
post_max_size=self.test_cfg["post_max_size"],
|
|||
|
|
),
|
|||
|
|
dtype=torch.long,
|
|||
|
|
device=boxes.device,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
boxes3d = boxes3d[keep]
|
|||
|
|
scores = scores[keep]
|
|||
|
|
labels = labels[keep]
|
|||
|
|
ret = dict(bboxes=boxes3d, scores=scores, labels=labels)
|
|||
|
|
ret_task.append(ret)
|
|||
|
|
rets.append(ret_task)
|
|||
|
|
else:
|
|||
|
|
rets.append(
|
|||
|
|
self.get_task_detections(
|
|||
|
|
num_class_with_bg,
|
|||
|
|
batch_cls_preds,
|
|||
|
|
batch_reg_preds,
|
|||
|
|
batch_cls_labels,
|
|||
|
|
metas,
|
|||
|
|
nms_scales[task_id],
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Merge branches results
|
|||
|
|
num_samples = len(rets[0])
|
|||
|
|
|
|||
|
|
ret_list = []
|
|||
|
|
for i in range(num_samples):
|
|||
|
|
for k in rets[0][i].keys():
|
|||
|
|
if k == "bboxes":
|
|||
|
|
bboxes = torch.cat([ret[i][k] for ret in rets])
|
|||
|
|
bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 5] * 0.5
|
|||
|
|
bboxes = metas[i]["box_type_3d"](bboxes, self.bbox_coder.code_size)
|
|||
|
|
elif k == "scores":
|
|||
|
|
scores = torch.cat([ret[i][k] for ret in rets])
|
|||
|
|
elif k == "labels":
|
|||
|
|
flag = 0
|
|||
|
|
for j, num_class in enumerate(self.num_classes):
|
|||
|
|
rets[j][i][k] += flag
|
|||
|
|
flag += num_class
|
|||
|
|
labels = torch.cat([ret[i][k].int() for ret in rets])
|
|||
|
|
ret_list.append([bboxes, scores, labels])
|
|||
|
|
return ret_list
|
|||
|
|
|
|||
|
|
def get_task_detections(
|
|||
|
|
self,
|
|||
|
|
num_class_with_bg,
|
|||
|
|
batch_cls_preds,
|
|||
|
|
batch_reg_preds,
|
|||
|
|
batch_cls_labels,
|
|||
|
|
metas,
|
|||
|
|
nms_scale=1.0,
|
|||
|
|
):
|
|||
|
|
"""Rotate nms for each task.
|
|||
|
|
Args:
|
|||
|
|
num_class_with_bg (int): Number of classes for the current task.
|
|||
|
|
batch_cls_preds (list[torch.Tensor]): Prediction score with the
|
|||
|
|
shape of [N].
|
|||
|
|
batch_reg_preds (list[torch.Tensor]): Prediction bbox with the
|
|||
|
|
shape of [N, 9].
|
|||
|
|
batch_cls_labels (list[torch.Tensor]): Prediction label with the
|
|||
|
|
shape of [N].
|
|||
|
|
metas (list[dict]): Meta information of each sample.
|
|||
|
|
Returns:
|
|||
|
|
list[dict[str: torch.Tensor]]: contains the following keys:
|
|||
|
|
-bboxes (torch.Tensor): Prediction bboxes after nms with the \
|
|||
|
|
shape of [N, 9].
|
|||
|
|
-scores (torch.Tensor): Prediction scores after nms with the \
|
|||
|
|
shape of [N].
|
|||
|
|
-labels (torch.Tensor): Prediction labels after nms with the \
|
|||
|
|
shape of [N].
|
|||
|
|
"""
|
|||
|
|
predictions_dicts = []
|
|||
|
|
post_center_range = self.test_cfg["post_center_limit_range"]
|
|||
|
|
if len(post_center_range) > 0:
|
|||
|
|
post_center_range = torch.tensor(
|
|||
|
|
post_center_range,
|
|||
|
|
dtype=batch_reg_preds[0].dtype,
|
|||
|
|
device=batch_reg_preds[0].device,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
for i, (box_preds, cls_preds, cls_labels) in enumerate(
|
|||
|
|
zip(batch_reg_preds, batch_cls_preds, batch_cls_labels)
|
|||
|
|
):
|
|||
|
|
|
|||
|
|
# Apply NMS in birdeye view
|
|||
|
|
|
|||
|
|
# get highest score per prediction, than apply nms
|
|||
|
|
# to remove overlapped box.
|
|||
|
|
if num_class_with_bg == 1:
|
|||
|
|
top_scores = cls_preds.squeeze(-1)
|
|||
|
|
top_labels = torch.zeros(
|
|||
|
|
cls_preds.shape[0], device=cls_preds.device, dtype=torch.long
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
top_labels = cls_labels.long()
|
|||
|
|
top_scores = cls_preds.squeeze(-1)
|
|||
|
|
|
|||
|
|
if self.test_cfg["score_threshold"] > 0.0:
|
|||
|
|
thresh = torch.tensor(
|
|||
|
|
[self.test_cfg["score_threshold"]], device=cls_preds.device
|
|||
|
|
).type_as(cls_preds)
|
|||
|
|
top_scores_keep = top_scores >= thresh
|
|||
|
|
top_scores = top_scores.masked_select(top_scores_keep)
|
|||
|
|
|
|||
|
|
if top_scores.shape[0] != 0:
|
|||
|
|
if self.test_cfg["score_threshold"] > 0.0:
|
|||
|
|
box_preds = box_preds[top_scores_keep]
|
|||
|
|
top_labels = top_labels[top_scores_keep]
|
|||
|
|
|
|||
|
|
bev_box = metas[i]["box_type_3d"](
|
|||
|
|
box_preds[:, :], self.bbox_coder.code_size
|
|||
|
|
).bev
|
|||
|
|
for cls, scale in enumerate(nms_scale):
|
|||
|
|
cur_bev_box = bev_box[top_labels == cls]
|
|||
|
|
cur_bev_box[:, [2, 3]] *= scale
|
|||
|
|
bev_box[top_labels == cls] = cur_bev_box
|
|||
|
|
boxes_for_nms = xywhr2xyxyr(bev_box)
|
|||
|
|
|
|||
|
|
# the nms in 3d detection just remove overlap boxes.
|
|||
|
|
|
|||
|
|
selected = nms_gpu(
|
|||
|
|
boxes_for_nms,
|
|||
|
|
top_scores,
|
|||
|
|
thresh=self.test_cfg["nms_thr"],
|
|||
|
|
pre_maxsize=self.test_cfg["pre_max_size"],
|
|||
|
|
post_max_size=self.test_cfg["post_max_size"],
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
selected = []
|
|||
|
|
|
|||
|
|
# if selected is not None:
|
|||
|
|
selected_boxes = box_preds[selected]
|
|||
|
|
selected_labels = top_labels[selected]
|
|||
|
|
selected_scores = top_scores[selected]
|
|||
|
|
|
|||
|
|
# finally generate predictions.
|
|||
|
|
if selected_boxes.shape[0] != 0:
|
|||
|
|
box_preds = selected_boxes
|
|||
|
|
scores = selected_scores
|
|||
|
|
label_preds = selected_labels
|
|||
|
|
final_box_preds = box_preds
|
|||
|
|
final_scores = scores
|
|||
|
|
final_labels = label_preds
|
|||
|
|
if post_center_range is not None:
|
|||
|
|
mask = (final_box_preds[:, :3] >= post_center_range[:3]).all(1)
|
|||
|
|
mask &= (final_box_preds[:, :3] <= post_center_range[3:]).all(1)
|
|||
|
|
predictions_dict = dict(
|
|||
|
|
bboxes=final_box_preds[mask],
|
|||
|
|
scores=final_scores[mask],
|
|||
|
|
labels=final_labels[mask],
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
predictions_dict = dict(
|
|||
|
|
bboxes=final_box_preds, scores=final_scores, labels=final_labels
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
dtype = batch_reg_preds[0].dtype
|
|||
|
|
device = batch_reg_preds[0].device
|
|||
|
|
predictions_dict = dict(
|
|||
|
|
bboxes=torch.zeros(
|
|||
|
|
[0, self.bbox_coder.code_size], dtype=dtype, device=device
|
|||
|
|
),
|
|||
|
|
scores=torch.zeros([0], dtype=dtype, device=device),
|
|||
|
|
labels=torch.zeros([0], dtype=top_labels.dtype, device=device),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
predictions_dicts.append(predictions_dict)
|
|||
|
|
return predictions_dicts
|