import torch
from mmcv.runner import force_fp32
from torch.nn import functional as F
from mmdet.models import DETECTORS
from .mvx_two_stage import MVXTwoStageDetector
[docs]@DETECTORS.register_module()
class MVXFasterRCNN(MVXTwoStageDetector):
"""Multi-modality VoxelNet using Faster R-CNN."""
def __init__(self, **kwargs):
super(MVXFasterRCNN, self).__init__(**kwargs)
[docs]@DETECTORS.register_module()
class DynamicMVXFasterRCNN(MVXTwoStageDetector):
"""Multi-modality VoxelNet using Faster R-CNN and dynamic voxelization."""
def __init__(self, **kwargs):
super(DynamicMVXFasterRCNN, self).__init__(**kwargs)
[docs] @torch.no_grad()
@force_fp32()
def voxelize(self, points):
"""Apply dynamic voxelization to points.
Args:
points (list[torch.Tensor]): Points of each sample.
Returns:
tuple[torch.Tensor]: Concatenated points and coordinates.
"""
coors = []
# dynamic voxelization only provide a coors mapping
for res in points:
res_coors = self.pts_voxel_layer(res)
coors.append(res_coors)
points = torch.cat(points, dim=0)
coors_batch = []
for i, coor in enumerate(coors):
coor_pad = F.pad(coor, (1, 0), mode='constant', value=i)
coors_batch.append(coor_pad)
coors_batch = torch.cat(coors_batch, dim=0)
return points, coors_batch