Skip to content

GitLab

Commit b5f1f52f authored by zhangwenwei's avatar zhangwenwei
Browse files

Merge branch 'refactor-voxel-encoder' into 'master' 

Clean detectors

See merge request open-mmlab/mmdet.3d!76
parents 90d0730e 10854be0
Show whitespace changes
Inline Side-by-side
Showing with 15 additions and 380 deletions
configs/_base_/models/pointpillars_second_fpn.py
View file @ b5f1f52f
......@@ -5,7 +5,7 @@
# keys in the config.
voxel_size = [0.25, 0.25, 8]
model = dict(
type='MVXFasterRCNNV2',
type='MVXFasterRCNN',
pts_voxel_layer=dict(
max_num_points=64,
point_cloud_range=[-50, -50, -5, 50, 50, 3],
......
configs/mvxnet/dv_mvx-v2_second_secfpn_fpn-fusion_adamw_2x8_80e_kitti-3d-3class.py
View file @ b5f1f52f
......@@ -3,7 +3,7 @@ voxel_size = [0.05, 0.05, 0.1]
point_cloud_range = [0, -40, -3, 70.4, 40, 1]
model = dict(
type='DynamicMVXFasterRCNNV2',
type='DynamicMVXFasterRCNN',
pretrained=('./pretrain_detectron/'
'ImageNetPretrained/MSRA/resnet50_msra.pth'),
img_backbone=dict(
......
configs/regnet/hv_pointpillars_regnet-400mf_fpn_sbn-all_4x8_2x_nus-3d.py
View file @ b5f1f52f
......@@ -6,7 +6,7 @@ _base_ = [
]
# model settings
model = dict(
type='MVXFasterRCNNV2',
type='MVXFasterRCNN',
pretrained=dict(pts='open-mmlab://regnetx_400mf'),
pts_backbone=dict(
_delete_=True,
......
mmdet3d/models/detectors/__init__.py
View file @ b5f1f52f
from .base import Base3DDetector
from .dynamic_voxelnet import DynamicVoxelNet
from .mvx_faster_rcnn import DynamicMVXFasterRCNN, DynamicMVXFasterRCNNV2
from .mvx_single_stage import MVXSingleStageDetector
from .mvx_faster_rcnn import DynamicMVXFasterRCNN, MVXFasterRCNN
from .mvx_two_stage import MVXTwoStageDetector
from .parta2 import PartA2
from .votenet import VoteNet
from .voxelnet import VoxelNet
__all__ = [
'Base3DDetector', 'VoxelNet', 'DynamicVoxelNet', 'MVXSingleStageDetector',
'MVXTwoStageDetector', 'DynamicMVXFasterRCNN', 'DynamicMVXFasterRCNNV2',
'PartA2', 'VoteNet'
'Base3DDetector', 'VoxelNet', 'DynamicVoxelNet', 'MVXTwoStageDetector',
'DynamicMVXFasterRCNN', 'MVXFasterRCNN', 'PartA2', 'VoteNet'
]
mmdet3d/models/detectors/mvx_faster_rcnn.py
View file @ b5f1f52f
......@@ -5,28 +5,19 @@ from mmdet.models import DETECTORS
from .mvx_two_stage import MVXTwoStageDetector
@DETECTORS.register_module()
class MVXFasterRCNN(MVXTwoStageDetector):
def __init__(self, **kwargs):
super(MVXFasterRCNN, self).__init__(**kwargs)
@DETECTORS.register_module()
class DynamicMVXFasterRCNN(MVXTwoStageDetector):
def __init__(self, **kwargs):
super(DynamicMVXFasterRCNN, self).__init__(**kwargs)
def extract_pts_feat(self, points, img_feats, img_metas):
if not self.with_pts_bbox:
return None
voxels, coors = self.voxelize(points)
# adopt an early fusion strategy
if self.with_fusion:
voxels = self.pts_fusion_layer(img_feats, points, voxels,
img_metas)
voxel_features, feature_coors = self.pts_voxel_encoder(voxels, coors)
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
@torch.no_grad()
def voxelize(self, points):
coors = []
......@@ -42,13 +33,6 @@ class DynamicMVXFasterRCNN(MVXTwoStageDetector):
coors_batch = torch.cat(coors_batch, dim=0)
return points, coors_batch
@DETECTORS.register_module()
class DynamicMVXFasterRCNNV2(DynamicMVXFasterRCNN):
def __init__(self, **kwargs):
super(DynamicMVXFasterRCNNV2, self).__init__(**kwargs)
def extract_pts_feat(self, points, img_feats, img_metas):
if not self.with_pts_bbox:
return None
......@@ -61,25 +45,3 @@ class DynamicMVXFasterRCNNV2(DynamicMVXFasterRCNN):
if self.with_pts_neck:
x = self.pts_neck(x)
return x
@DETECTORS.register_module()
class MVXFasterRCNNV2(MVXTwoStageDetector):
def __init__(self, **kwargs):
super(MVXFasterRCNNV2, self).__init__(**kwargs)
def extract_pts_feat(self, pts, img_feats, img_metas):
if not self.with_pts_bbox:
return None
voxels, num_points, coors = self.voxelize(pts)
voxel_features = self.pts_voxel_encoder(voxels, num_points, coors,
img_feats, img_metas)
batch_size = coors[-1, 0] + 1
x = self.pts_middle_encoder(voxel_features, coors, batch_size)
x = self.pts_backbone(x)
if self.with_pts_neck:
x = self.pts_neck(x)
return x
mmdet3d/models/detectors/mvx_single_stage.py deleted 100644 → 0
View file @ 90d0730e
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmdet3d.core import bbox3d2result
from mmdet3d.ops import Voxelization
from mmdet.models import DETECTORS
from .. import builder
from .single_stage import SingleStage3DDetector
@DETECTORS.register_module()
class MVXSingleStageDetector(SingleStage3DDetector):
def __init__(self,
voxel_layer,
voxel_encoder,
middle_encoder,
fusion_layer,
img_backbone,
pts_backbone,
img_neck=None,
pts_neck=None,
pts_bbox_head=None,
img_bbox_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(MVXSingleStageDetector, self).__init__()
self.voxel_layer = Voxelization(**voxel_layer)
self.voxel_encoder = builder.build_voxel_encoder(voxel_encoder)
self.middle_encoder = builder.build_middle_encoder(middle_encoder)
self.pts_backbone = builder.build_backbone(pts_backbone)
if fusion_layer:
self.fusion_layer = builder.build_fusion_layer(fusion_layer)
if img_backbone:
self.img_backbone = builder.build_backbone(img_backbone)
pts_bbox_head.update(train_cfg=train_cfg)
pts_bbox_head.update(test_cfg=test_cfg)
self.pts_bbox_head = builder.build_head(pts_bbox_head)
if img_neck is not None:
self.img_neck = builder.build_neck(img_neck)
if pts_neck is not None:
self.pts_neck = builder.build_neck(pts_neck)
if img_bbox_head is not None:
self.img_bbox_head = builder.build_head(img_bbox_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def init_weights(self, pretrained=None):
super(MVXSingleStageDetector, self).init_weights(pretrained)
if self.with_img_backbone:
self.img_backbone.init_weights(pretrained=pretrained)
if self.with_img_neck:
if isinstance(self.img_neck, nn.Sequential):
for m in self.img_neck:
m.init_weights()
else:
self.img_neck.init_weights()
if self.with_img_bbox:
self.img_bbox_head.init_weights()
if self.with_pts_bbox:
self.pts_bbox_head.init_weights()
@property
def with_pts_bbox(self):
return hasattr(self,
'pts_bbox_head') and self.pts_bbox_head is not None
@property
def with_img_bbox(self):
return hasattr(self,
'img_bbox_head') and self.img_bbox_head is not None
@property
def with_img_backbone(self):
return hasattr(self, 'img_backbone') and self.img_backbone is not None
@property
def with_fusion(self):
return hasattr(self, 'fusion_layer') and self.fusion_layer is not None
@property
def with_img_neck(self):
return hasattr(self, 'img_neck') and self.img_neck is not None
@property
def with_pts_neck(self):
return hasattr(self, 'pts_neck') and self.pts_neck is not None
def extract_feat(self, points, img, img_metas):
if self.with_img_backbone:
img_feats = self.img_backbone(img)
if self.with_img_neck:
img_feats = self.img_neck(img_feats)
voxels, num_points, coors = self.voxelize(points)
voxel_features = self.voxel_encoder(voxels, num_points, coors)
batch_size = coors[-1, 0] + 1
x = self.middle_encoder(voxel_features, coors, batch_size)
x = self.pts_backbone(x)
if self.with_neck:
x = self.pts_neck(x)
return x
@torch.no_grad()
def voxelize(self, points):
voxels, coors, num_points = [], [], []
for res in points:
res_voxels, res_coors, res_num_points = self.voxel_layer(res)
voxels.append(res_voxels)
coors.append(res_coors)
num_points.append(res_num_points)
voxels = torch.cat(voxels, dim=0)
num_points = torch.cat(num_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 voxels, num_points, coors_batch
def forward_train(self,
points,
img_metas,
gt_bboxes_3d,
gt_labels,
img=None,
gt_bboxes_ignore=None):
x = self.extract_feat(points, img=img, img_metas=img_metas)
outs = self.pts_bbox_head(x)
loss_inputs = outs + (gt_bboxes_3d, gt_labels, img_metas)
losses = self.pts_bbox_head.loss(
*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
return losses
def simple_test(self,
points,
img_metas,
img=None,
gt_bboxes_3d=None,
rescale=False):
x = self.extract_feat(points, img, img_metas)
outs = self.pts_bbox_head(x)
bbox_list = self.pts_bbox_head.get_bboxes(
*outs, img_metas, rescale=rescale)
bbox_results = [
bbox3d2result(bboxes, scores, labels)
for bboxes, scores, labels in bbox_list
]
return bbox_results[0]
def aug_test(self, points, imgs, img_metas, rescale=False):
raise NotImplementedError
@DETECTORS.register_module()
class DynamicMVXNet(MVXSingleStageDetector):
def __init__(self,
voxel_layer,
voxel_encoder,
middle_encoder,
pts_backbone,
fusion_layer=None,
img_backbone=None,
img_neck=None,
pts_neck=None,
pts_bbox_head=None,
img_bbox_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(DynamicMVXNet, self).__init__(
voxel_layer=voxel_layer,
voxel_encoder=voxel_encoder,
middle_encoder=middle_encoder,
img_backbone=img_backbone,
fusion_layer=fusion_layer,
pts_backbone=pts_backbone,
pts_neck=pts_neck,
img_neck=img_neck,
img_bbox_head=img_bbox_head,
pts_bbox_head=pts_bbox_head,
train_cfg=train_cfg,
test_cfg=test_cfg,
pretrained=pretrained,
)
def extract_feat(self, points, img, img_metas):
if self.with_img_backbone:
img_feats = self.img_backbone(img)
if self.with_img_neck:
img_feats = self.img_neck(img_feats)
voxels, coors = self.voxelize(points)
# adopt an early fusion strategy
if self.with_fusion:
voxels = self.fusion_layer(img_feats, points, voxels, img_metas)
voxel_features, feature_coors = self.voxel_encoder(voxels, coors)
batch_size = coors[-1, 0] + 1
x = self.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
@torch.no_grad()
def voxelize(self, points):
coors = []
# dynamic voxelization only provide a coors mapping
for res in points:
res_coors = self.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
@DETECTORS.register_module()
class DynamicMVXNetV2(DynamicMVXNet):
def __init__(self,
voxel_layer,
voxel_encoder,
middle_encoder,
pts_backbone,
fusion_layer=None,
img_backbone=None,
img_neck=None,
pts_neck=None,
pts_bbox_head=None,
img_bbox_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(DynamicMVXNetV2, self).__init__(
voxel_layer=voxel_layer,
voxel_encoder=voxel_encoder,
middle_encoder=middle_encoder,
img_backbone=img_backbone,
fusion_layer=fusion_layer,
pts_backbone=pts_backbone,
pts_neck=pts_neck,
img_neck=img_neck,
img_bbox_head=img_bbox_head,
pts_bbox_head=pts_bbox_head,
train_cfg=train_cfg,
test_cfg=test_cfg,
pretrained=pretrained,
)
def extract_feat(self, points, img, img_metas):
if self.with_img_backbone:
img_feats = self.img_backbone(img)
if self.with_img_neck:
img_feats = self.img_neck(img_feats)
voxels, coors = self.voxelize(points)
voxel_features, feature_coors = self.voxel_encoder(
voxels, coors, points, img_feats, img_metas)
batch_size = coors[-1, 0] + 1
x = self.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
@DETECTORS.register_module()
class DynamicMVXNetV3(DynamicMVXNet):
def __init__(self,
voxel_layer,
voxel_encoder,
middle_encoder,
pts_backbone,
fusion_layer=None,
img_backbone=None,
img_neck=None,
pts_neck=None,
pts_bbox_head=None,
img_bbox_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(DynamicMVXNetV3, self).__init__(
voxel_layer=voxel_layer,
voxel_encoder=voxel_encoder,
middle_encoder=middle_encoder,
img_backbone=img_backbone,
fusion_layer=fusion_layer,
pts_backbone=pts_backbone,
pts_neck=pts_neck,
img_neck=img_neck,
img_bbox_head=img_bbox_head,
pts_bbox_head=pts_bbox_head,
train_cfg=train_cfg,
test_cfg=test_cfg,
pretrained=pretrained,
)
def extract_feat(self, points, img, img_metas):
if self.with_img_backbone:
img_feats = self.img_backbone(img)
if self.with_img_neck:
img_feats = self.img_neck(img_feats)
voxels, coors = self.voxelize(points)
voxel_features, feature_coors = self.voxel_encoder(voxels, coors)
batch_size = coors[-1, 0] + 1
x = self.middle_encoder(voxel_features, feature_coors, batch_size)
x = self.pts_backbone(x)
if self.with_pts_neck:
x = self.pts_neck(x, coors, points, img_feats, img_metas)
return x
mmdet3d/models/detectors/mvx_two_stage.py
View file @ b5f1f52f
......@@ -165,7 +165,8 @@ class MVXTwoStageDetector(Base3DDetector):
if not self.with_pts_bbox:
return None
voxels, num_points, coors = self.voxelize(pts)
voxel_features = self.pts_voxel_encoder(voxels, num_points, coors)
voxel_features = self.pts_voxel_encoder(voxels, num_points, coors,
img_feats, img_metas)
batch_size = coors[-1, 0] + 1
x = self.pts_middle_encoder(voxel_features, coors, batch_size)
x = self.pts_backbone(x)
......
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