Source code for mmdet3d.models.detectors.mvx_faster_rcnn
# Copyright (c) OpenMMLab. All rights reserved.
import torch
from mmcv.runner import force_fp32
from torch.nn import functional as F
from ..builder 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
[docs] def extract_pts_feat(self, points, img_feats, img_metas):
"""Extract point features."""
if not self.with_pts_bbox:
return None
voxels, coors = self.voxelize(points)
voxel_features, feature_coors = self.pts_voxel_encoder(
voxels, coors, points, img_feats, img_metas)
batch_size = coors[-1, 0] + 1
x = self.pts_middle_encoder(voxel_features, feature_coors, batch_size)
x = self.pts_backbone(x)
if self.with_pts_neck:
x = self.pts_neck(x)
return x