Skip to content

GitLab

Commit 4f1a5e52 authored by liyinhao's avatar liyinhao
Browse files

Merge branch 'master_temp' into indoor_augment

parents c2c0f3d8 f584b970
Hide whitespace changes
Inline Side-by-side
Showing with 860 additions and 272 deletions
mmdet3d/models/middle_encoders/sparse_encoder.py
View file @ 4f1a5e52
import torch.nn as nn
from mmcv.cnn import build_norm_layer
import mmdet3d.ops.spconv as spconv
from mmdet.ops import build_norm_layer
from ..registry import MIDDLE_ENCODERS
@MIDDLE_ENCODERS.register_module
@MIDDLE_ENCODERS.register_module()
class SparseEncoder(nn.Module):
def __init__(self,
......
mmdet3d/models/middle_encoders/sparse_unet.py 0 → 100644
View file @ 4f1a5e52
import torch
import torch.nn as nn
from mmcv.cnn import build_norm_layer
import mmdet3d.ops.spconv as spconv
from mmdet3d.ops import SparseBasicBlock
from ..registry import MIDDLE_ENCODERS
@MIDDLE_ENCODERS.register_module()
class SparseUNet(nn.Module):
def __init__(self,
in_channels,
output_shape,
pre_act=False,
norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
base_channels=16,
output_channels=128,
encoder_channels=((16, ), (32, 32, 32), (64, 64, 64), (64, 64,
64)),
encoder_paddings=((1, ), (1, 1, 1), (1, 1, 1), ((0, 1, 1), 1,
1)),
decoder_channels=((64, 64, 64), (64, 64, 32), (32, 32, 16),
(16, 16, 16)),
decoder_paddings=((1, 0), (1, 0), (0, 0), (0, 1))):
"""SparseUNet for PartA^2
See https://arxiv.org/abs/1907.03670 for more detials.
Args:
in_channels (int): the number of input channels
output_shape (list[int]): the shape of output tensor
pre_act (bool): use pre_act_block or post_act_block
norm_cfg (dict): config of normalization layer
base_channels (int): out channels for conv_input layer
output_channels (int): out channels for conv_out layer
encoder_channels (tuple[tuple[int]]):
conv channels of each encode block
encoder_paddings (tuple[tuple[int]]): paddings of each encode block
decoder_channels (tuple[tuple[int]]):
conv channels of each decode block
decoder_paddings (tuple[tuple[int]]): paddings of each decode block
"""
super().__init__()
self.sparse_shape = output_shape
self.output_shape = output_shape
self.in_channels = in_channels
self.pre_act = pre_act
self.base_channels = base_channels
self.output_channels = output_channels
self.encoder_channels = encoder_channels
self.encoder_paddings = encoder_paddings
self.decoder_channels = decoder_channels
self.decoder_paddings = decoder_paddings
self.stage_num = len(self.encoder_channels)
# Spconv init all weight on its own
if pre_act:
# TODO: use ConvModule to encapsulate
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(
in_channels,
self.base_channels,
3,
padding=1,
bias=False,
indice_key='subm1'))
make_block = self.pre_act_block
else:
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(
in_channels,
self.base_channels,
3,
padding=1,
bias=False,
indice_key='subm1'),
build_norm_layer(norm_cfg, self.base_channels)[1], nn.ReLU())
make_block = self.post_act_block
encoder_out_channels = self.make_encoder_layers(
make_block, norm_cfg, self.base_channels)
self.make_decoder_layers(make_block, norm_cfg, encoder_out_channels)
self.conv_out = spconv.SparseSequential(
# [200, 176, 5] -> [200, 176, 2]
spconv.SparseConv3d(
encoder_out_channels,
self.output_channels, (3, 1, 1),
stride=(2, 1, 1),
padding=0,
bias=False,
indice_key='spconv_down2'),
build_norm_layer(norm_cfg, self.output_channels)[1],
nn.ReLU())
def forward(self, voxel_features, coors, batch_size):
"""Forward of SparseUNet
Args:
voxel_features (torch.float32): shape [N, C]
coors (torch.int32): shape [N, 4](batch_idx, z_idx, y_idx, x_idx)
batch_size (int): batch size
Returns:
dict: backbone features
"""
coors = coors.int()
input_sp_tensor = spconv.SparseConvTensor(voxel_features, coors,
self.sparse_shape,
batch_size)
x = self.conv_input(input_sp_tensor)
encode_features = []
for encoder_layer in self.encoder_layers:
x = encoder_layer(x)
encode_features.append(x)
# for detection head
# [200, 176, 5] -> [200, 176, 2]
out = self.conv_out(encode_features[-1])
spatial_features = out.dense()
N, C, D, H, W = spatial_features.shape
spatial_features = spatial_features.view(N, C * D, H, W)
# for segmentation head, with output shape:
# [400, 352, 11] <- [200, 176, 5]
# [800, 704, 21] <- [400, 352, 11]
# [1600, 1408, 41] <- [800, 704, 21]
# [1600, 1408, 41] <- [1600, 1408, 41]
decode_features = []
x = encode_features[-1]
for i in range(self.stage_num, 0, -1):
x = self.decoder_layer_forward(encode_features[i - 1], x,
getattr(self, f'lateral_layer{i}'),
getattr(self, f'merge_layer{i}'),
getattr(self, f'upsample_layer{i}'))
decode_features.append(x)
seg_features = decode_features[-1].features
ret = dict(
spatial_features=spatial_features, seg_features=seg_features)
return ret
def decoder_layer_forward(self, x_lateral, x_bottom, lateral_layer,
merge_layer, upsample_layer):
"""Forward of upsample and residual block.
Args:
x_lateral (SparseConvTensor): lateral tensor
x_bottom (SparseConvTensor): feature from bottom layer
lateral_layer (SparseBasicBlock): convolution for lateral tensor
merge_layer (SparseSequential): convolution for merging features
upsample_layer (SparseSequential): convolution for upsampling
Returns:
SparseConvTensor: upsampled feature
"""
x = lateral_layer(x_lateral)
x.features = torch.cat((x_bottom.features, x.features), dim=1)
x_merge = merge_layer(x)
x = self.reduce_channel(x, x_merge.features.shape[1])
x.features = x_merge.features + x.features
x = upsample_layer(x)
return x
@staticmethod
def reduce_channel(x, out_channels):
"""reduce channel for element-wise addition.
Args:
x (SparseConvTensor): x.features (N, C1)
out_channels (int): the number of channel after reduction
Returns:
SparseConvTensor: channel reduced feature
"""
features = x.features
n, in_channels = features.shape
assert (in_channels % out_channels
== 0) and (in_channels >= out_channels)
x.features = features.view(n, out_channels, -1).sum(dim=2)
return x
def pre_act_block(self,
in_channels,
out_channels,
kernel_size,
indice_key=None,
stride=1,
padding=0,
conv_type='subm',
norm_cfg=None):
"""Make pre activate sparse convolution block.
Args:
in_channels (int): the number of input channels
out_channels (int): the number of out channels
kernel_size (int): kernel size of convolution
indice_key (str): the indice key used for sparse tensor
stride (int): the stride of convolution
padding (int or list[int]): the padding number of input
conv_type (str): conv type in 'subm', 'spconv' or 'inverseconv'
norm_cfg (dict): config of normalization layer
Returns:
spconv.SparseSequential: pre activate sparse convolution block.
"""
# TODO: use ConvModule to encapsulate
assert conv_type in ['subm', 'spconv', 'inverseconv']
if conv_type == 'subm':
m = spconv.SparseSequential(
build_norm_layer(norm_cfg, in_channels)[1],
nn.ReLU(inplace=True),
spconv.SubMConv3d(
in_channels,
out_channels,
kernel_size,
padding=padding,
bias=False,
indice_key=indice_key))
elif conv_type == 'spconv':
m = spconv.SparseSequential(
build_norm_layer(norm_cfg, in_channels)[1],
nn.ReLU(inplace=True),
spconv.SparseConv3d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
bias=False,
indice_key=indice_key))
elif conv_type == 'inverseconv':
m = spconv.SparseSequential(
build_norm_layer(norm_cfg, in_channels)[1],
nn.ReLU(inplace=True),
spconv.SparseInverseConv3d(
in_channels,
out_channels,
kernel_size,
bias=False,
indice_key=indice_key))
else:
raise NotImplementedError
return m
def post_act_block(self,
in_channels,
out_channels,
kernel_size,
indice_key,
stride=1,
padding=0,
conv_type='subm',
norm_cfg=None):
"""Make post activate sparse convolution block.
Args:
in_channels (int): the number of input channels
out_channels (int): the number of out channels
kernel_size (int): kernel size of convolution
indice_key (str): the indice key used for sparse tensor
stride (int): the stride of convolution
padding (int or list[int]): the padding number of input
conv_type (str): conv type in 'subm', 'spconv' or 'inverseconv'
norm_cfg (dict[str]): config of normalization layer
Returns:
spconv.SparseSequential: post activate sparse convolution block.
"""
# TODO: use ConvModule to encapsulate
assert conv_type in ['subm', 'spconv', 'inverseconv']
if conv_type == 'subm':
m = spconv.SparseSequential(
spconv.SubMConv3d(
in_channels,
out_channels,
kernel_size,
bias=False,
indice_key=indice_key),
build_norm_layer(norm_cfg, out_channels)[1],
nn.ReLU(inplace=True))
elif conv_type == 'spconv':
m = spconv.SparseSequential(
spconv.SparseConv3d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
bias=False,
indice_key=indice_key),
build_norm_layer(norm_cfg, out_channels)[1],
nn.ReLU(inplace=True))
elif conv_type == 'inverseconv':
m = spconv.SparseSequential(
spconv.SparseInverseConv3d(
in_channels,
out_channels,
kernel_size,
bias=False,
indice_key=indice_key),
build_norm_layer(norm_cfg, out_channels)[1],
nn.ReLU(inplace=True))
else:
raise NotImplementedError
return m
def make_encoder_layers(self, make_block, norm_cfg, in_channels):
"""make encoder layers using sparse convs
Args:
make_block (method): a bounded function to build blocks
norm_cfg (dict[str]): config of normalization layer
in_channels (int): the number of encoder input channels
Returns:
int: the number of encoder output channels
"""
self.encoder_layers = spconv.SparseSequential()
for i, blocks in enumerate(self.encoder_channels):
blocks_list = []
for j, out_channels in enumerate(tuple(blocks)):
padding = tuple(self.encoder_paddings[i])[j]
# each stage started with a spconv layer
# except the first stage
if i != 0 and j == 0:
blocks_list.append(
make_block(
in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
stride=2,
padding=padding,
indice_key=f'spconv{i + 1}',
conv_type='spconv'))
else:
blocks_list.append(
make_block(
in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
padding=padding,
indice_key=f'subm{i + 1}'))
in_channels = out_channels
stage_name = f'encoder_layer{i + 1}'
stage_layers = spconv.SparseSequential(*blocks_list)
self.encoder_layers.add_module(stage_name, stage_layers)
return out_channels
def make_decoder_layers(self, make_block, norm_cfg, in_channels):
"""make decoder layers using sparse convs
Args:
make_block (method): a bounded function to build blocks
norm_cfg (dict[str]): config of normalization layer
in_channels (int): the number of encoder input channels
Returns:
int: the number of encoder output channels
"""
block_num = len(self.decoder_channels)
for i, block_channels in enumerate(self.decoder_channels):
paddings = self.decoder_paddings[i]
setattr(
self, f'lateral_layer{block_num - i}',
SparseBasicBlock(
in_channels,
block_channels[0],
conv_cfg=dict(
type='SubMConv3d', indice_key=f'subm{block_num - i}'),
norm_cfg=norm_cfg))
setattr(
self, f'merge_layer{block_num - i}',
make_block(
in_channels * 2,
block_channels[1],
3,
norm_cfg=norm_cfg,
padding=paddings[0],
indice_key=f'subm{block_num - i}'))
if block_num - i != 1:
setattr(
self, f'upsample_layer{block_num - i}',
make_block(
in_channels,
block_channels[2],
3,
norm_cfg=norm_cfg,
padding=paddings[1],
indice_key=f'spconv{block_num - i}',
conv_type='inverseconv'))
else:
# use submanifold conv instead of inverse conv
# in the last block
setattr(
self, f'upsample_layer{block_num - i}',
make_block(
in_channels,
block_channels[2],
3,
norm_cfg=norm_cfg,
padding=paddings[1],
indice_key='subm1',
conv_type='subm'))
in_channels = block_channels[2]
mmdet3d/models/necks/second_fpn.py
View file @ 4f1a5e52
......@@ -2,16 +2,15 @@ from functools import partial
import torch
import torch.nn as nn
from mmcv.cnn import constant_init, kaiming_init
from mmcv.cnn import build_norm_layer, constant_init, kaiming_init
from torch.nn import Sequential
from torch.nn.modules.batchnorm import _BatchNorm
from mmdet.models import NECKS
from mmdet.ops import build_norm_layer
from .. import builder
@NECKS.register_module
@NECKS.register_module()
class SECONDFPN(nn.Module):
"""Compare with RPN, RPNV2 support arbitrary number of stage.
"""
......@@ -64,7 +63,7 @@ class SECONDFPN(nn.Module):
return [out]
@NECKS.register_module
@NECKS.register_module()
class SECONDFusionFPN(SECONDFPN):
"""Compare with RPN, RPNV2 support arbitrary number of stage.
"""
......
mmdet3d/models/registry.py
View file @ 4f1a5e52
from mmdet.utils import Registry
from mmcv.utils import Registry
VOXEL_ENCODERS = Registry('voxel_encoder')
MIDDLE_ENCODERS = Registry('middle_encoder')
......
mmdet3d/models/roi_heads/__init__.py 0 → 100644
View file @ 4f1a5e52
from .mask_heads import PointwiseSemanticHead
__all__ = ['PointwiseSemanticHead']
mmdet3d/models/bbox_heads/__init__.py → mmdet3d/models/roi_heads/bbox_heads/__init__.py
View file @ 4f1a5e52
from mmdet.models.bbox_heads import (BBoxHead, ConvFCBBoxHead,
DoubleConvFCBBoxHead, Shared2FCBBoxHead,
Shared4Conv1FCBBoxHead)
from mmdet.models.roi_heads.bbox_heads import (BBoxHead, ConvFCBBoxHead,
DoubleConvFCBBoxHead,
Shared2FCBBoxHead,
Shared4Conv1FCBBoxHead)
__all__ = [
'BBoxHead', 'ConvFCBBoxHead', 'Shared2FCBBoxHead',
......
mmdet3d/models/roi_heads/mask_heads/__init__.py 0 → 100644
View file @ 4f1a5e52
from .pointwise_semantic_head import PointwiseSemanticHead
__all__ = ['PointwiseSemanticHead']
mmdet3d/models/roi_heads/mask_heads/pointwise_semantic_head.py 0 → 100644
View file @ 4f1a5e52
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmdet3d.core import multi_apply
from mmdet3d.core.bbox import box_torch_ops
from mmdet3d.models.builder import build_loss
from mmdet3d.ops.roiaware_pool3d import points_in_boxes_gpu
from mmdet.models import HEADS
@HEADS.register_module()
class PointwiseSemanticHead(nn.Module):
"""Semantic segmentation head for point-wise segmentation.
Predict point-wise segmentation and part regression results for PartA2.
See https://arxiv.org/abs/1907.03670 for more detials.
Args:
in_channels (int): the number of input channel.
num_classes (int): the number of class.
extra_width (float): boxes enlarge width.
loss_seg (dict): Config of segmentation loss.
loss_part (dict): Config of part prediction loss.
"""
def __init__(self,
in_channels,
num_classes=3,
extra_width=0.2,
seg_score_thr=0.3,
loss_seg=dict(
type='FocalLoss',
use_sigmoid=True,
reduction='sum',
gamma=2.0,
alpha=0.25,
loss_weight=1.0),
loss_part=dict(
type='CrossEntropyLoss',
use_sigmoid=True,
loss_weight=1.0)):
super(PointwiseSemanticHead, self).__init__()
self.extra_width = extra_width
self.num_classes = num_classes
self.seg_score_thr = seg_score_thr
self.seg_cls_layer = nn.Linear(in_channels, 1, bias=True)
self.seg_reg_layer = nn.Linear(in_channels, 3, bias=True)
self.loss_seg = build_loss(loss_seg)
self.loss_part = build_loss(loss_part)
def forward(self, x):
seg_preds = self.seg_cls_layer(x) # (N, 1)
part_preds = self.seg_reg_layer(x) # (N, 3)
seg_scores = torch.sigmoid(seg_preds).detach()
seg_mask = (seg_scores > self.seg_score_thr)
part_offsets = torch.sigmoid(part_preds).clone().detach()
part_offsets[seg_mask.view(-1) == 0] = 0
part_feats = torch.cat((part_offsets, seg_scores),
dim=-1) # shape (npoints, 4)
return dict(
seg_preds=seg_preds, part_preds=part_preds, part_feats=part_feats)
def get_targets_single(self, voxel_centers, gt_bboxes_3d, gt_labels_3d):
"""generate segmentation and part prediction targets
Args:
voxel_centers (torch.Tensor): shape [voxel_num, 3],
the center of voxels
gt_bboxes_3d (torch.Tensor): shape [box_num, 7], gt boxes
gt_labels_3d (torch.Tensor): shape [box_num], class label of gt
Returns:
tuple : segmentation targets with shape [voxel_num]
part prediction targets with shape [voxel_num, 3]
"""
enlarged_gt_boxes = box_torch_ops.enlarge_box3d_lidar(
gt_bboxes_3d, extra_width=self.extra_width)
part_targets = voxel_centers.new_zeros((voxel_centers.shape[0], 3),
dtype=torch.float32)
box_idx = points_in_boxes_gpu(
voxel_centers.unsqueeze(0),
gt_bboxes_3d.unsqueeze(0)).squeeze(0) # -1 ~ box_num
enlarge_box_idx = points_in_boxes_gpu(
voxel_centers.unsqueeze(0),
enlarged_gt_boxes.unsqueeze(0)).squeeze(0).long() # -1 ~ box_num
gt_labels_pad = F.pad(
gt_labels_3d, (1, 0), mode='constant', value=self.num_classes)
seg_targets = gt_labels_pad[(box_idx.long() + 1)]
fg_pt_flag = box_idx > -1
ignore_flag = fg_pt_flag ^ (enlarge_box_idx > -1)
seg_targets[ignore_flag] = -1
for k in range(gt_bboxes_3d.shape[0]):
k_box_flag = box_idx == k
# no point in current box (caused by velodyne reduce)
if not k_box_flag.any():
continue
fg_voxels = voxel_centers[k_box_flag]
transformed_voxels = fg_voxels - gt_bboxes_3d[k, 0:3]
transformed_voxels = box_torch_ops.rotation_3d_in_axis(
transformed_voxels.unsqueeze(0),
-gt_bboxes_3d[k, 6].view(1),
axis=2)
part_targets[k_box_flag] = transformed_voxels / gt_bboxes_3d[
k, 3:6] + voxel_centers.new_tensor([0.5, 0.5, 0])
part_targets = torch.clamp(part_targets, min=0)
return seg_targets, part_targets
def get_targets(self, voxels_dict, gt_bboxes_3d, gt_labels_3d):
batch_size = len(gt_labels_3d)
voxel_center_list = []
for idx in range(batch_size):
coords_idx = voxels_dict['coors'][:, 0] == idx
voxel_center_list.append(voxels_dict['voxel_centers'][coords_idx])
seg_targets, part_targets = multi_apply(self.get_targets_single,
voxel_center_list,
gt_bboxes_3d, gt_labels_3d)
seg_targets = torch.cat(seg_targets, dim=0)
part_targets = torch.cat(part_targets, dim=0)
return dict(seg_targets=seg_targets, part_targets=part_targets)
def loss(self, seg_preds, part_preds, seg_targets, part_targets):
"""Calculate point-wise segmentation and part prediction losses.
Args:
seg_preds (torch.Tensor): prediction of binary
segmentation with shape [voxel_num, 1].
part_preds (torch.Tensor): prediction of part
with shape [voxel_num, 3].
seg_targets (torch.Tensor): target of segmentation
with shape [voxel_num, 1].
part_targets (torch.Tensor): target of part with
shape [voxel_num, 3].
Returns:
dict: loss of segmentation and part prediction.
"""
pos_mask = (seg_targets > -1) & (seg_targets < self.num_classes)
binary_seg_target = pos_mask.long()
pos = pos_mask.float()
neg = (seg_targets == self.num_classes).float()
seg_weights = pos + neg
pos_normalizer = pos.sum()
seg_weights = seg_weights / torch.clamp(pos_normalizer, min=1.0)
loss_seg = self.loss_seg(seg_preds, binary_seg_target, seg_weights)
if pos_normalizer > 0:
loss_part = self.loss_part(part_preds[pos_mask],
part_targets[pos_mask])
else:
# fake a part loss
loss_part = loss_seg.new_tensor(0)
return dict(loss_seg=loss_seg, loss_part=loss_part)
mmdet3d/models/roi_extractors/__init__.py → mmdet3d/models/roi_heads/roi_extractors/__init__.py
View file @ 4f1a5e52
from mmdet.models.roi_extractors.single_level import SingleRoIExtractor
from mmdet.models.roi_heads.roi_extractors import SingleRoIExtractor
__all__ = ['SingleRoIExtractor']
mmdet3d/models/utils/__init__.py deleted 100644 → 0
View file @ c2c0f3d8
from mmdet.models.utils import ResLayer, bias_init_with_prob
__all__ = ['bias_init_with_prob', 'ResLayer']
mmdet3d/models/utils/weight_init.py deleted 100644 → 0
View file @ c2c0f3d8
import numpy as np
import torch.nn as nn
def xavier_init(module, gain=1, bias=0, distribution='normal'):
assert distribution in ['uniform', 'normal']
if distribution == 'uniform':
nn.init.xavier_uniform_(module.weight, gain=gain)
else:
nn.init.xavier_normal_(module.weight, gain=gain)
if hasattr(module, 'bias'):
nn.init.constant_(module.bias, bias)
def normal_init(module, mean=0, std=1, bias=0):
nn.init.normal_(module.weight, mean, std)
if hasattr(module, 'bias'):
nn.init.constant_(module.bias, bias)
def uniform_init(module, a=0, b=1, bias=0):
nn.init.uniform_(module.weight, a, b)
if hasattr(module, 'bias'):
nn.init.constant_(module.bias, bias)
def kaiming_init(module,
mode='fan_out',
nonlinearity='relu',
bias=0,
distribution='normal'):
assert distribution in ['uniform', 'normal']
if distribution == 'uniform':
nn.init.kaiming_uniform_(
module.weight, mode=mode, nonlinearity=nonlinearity)
else:
nn.init.kaiming_normal_(
module.weight, mode=mode, nonlinearity=nonlinearity)
if hasattr(module, 'bias'):
nn.init.constant_(module.bias, bias)
def bias_init_with_prob(prior_prob):
""" initialize conv/fc bias value according to giving probablity"""
bias_init = float(-np.log((1 - prior_prob) / prior_prob))
return bias_init
mmdet3d/models/voxel_encoders/pillar_encoder.py
View file @ 4f1a5e52
import torch
from mmcv.cnn import build_norm_layer
from torch import nn
from mmdet3d.ops import DynamicScatter
from mmdet.ops import build_norm_layer
from ..registry import VOXEL_ENCODERS
from .utils import PFNLayer, get_paddings_indicator
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class PillarFeatureNet(nn.Module):
def __init__(self,
......@@ -118,7 +118,7 @@ class PillarFeatureNet(nn.Module):
return features.squeeze()
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class DynamicPillarFeatureNet(PillarFeatureNet):
def __init__(self,
......@@ -237,7 +237,7 @@ class DynamicPillarFeatureNet(PillarFeatureNet):
return voxel_feats, voxel_coors
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class AlignedPillarFeatureNet(nn.Module):
def __init__(self,
......
mmdet3d/models/voxel_encoders/utils.py
View file @ 4f1a5e52
import torch
from mmcv.cnn import build_norm_layer
from torch import nn
from torch.nn import functional as F
from mmdet.ops import build_norm_layer
class Empty(nn.Module):
......
mmdet3d/models/voxel_encoders/voxel_encoder.py
View file @ 4f1a5e52
import torch
from mmcv.cnn import build_norm_layer
from torch import nn
from torch.nn import functional as F
from mmdet3d.ops import DynamicScatter
from mmdet.ops import build_norm_layer
from .. import builder
from ..registry import VOXEL_ENCODERS
from .utils import Empty, VFELayer, get_paddings_indicator
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class VoxelFeatureExtractor(nn.Module):
def __init__(self,
......@@ -71,7 +71,7 @@ class VoxelFeatureExtractor(nn.Module):
return voxelwise
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class VoxelFeatureExtractorV2(nn.Module):
def __init__(self,
......@@ -132,7 +132,7 @@ class VoxelFeatureExtractorV2(nn.Module):
return voxelwise
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class VoxelFeatureExtractorV3(nn.Module):
def __init__(self,
......@@ -152,7 +152,7 @@ class VoxelFeatureExtractorV3(nn.Module):
return points_mean.contiguous()
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class DynamicVFEV3(nn.Module):
def __init__(self,
......@@ -170,7 +170,7 @@ class DynamicVFEV3(nn.Module):
return features, features_coors
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class DynamicVFE(nn.Module):
def __init__(self,
......@@ -318,7 +318,7 @@ class DynamicVFE(nn.Module):
return voxel_feats, voxel_coors
@VOXEL_ENCODERS.register_module
@VOXEL_ENCODERS.register_module()
class HardVFE(nn.Module):
def __init__(self,
......
mmdet3d/ops/__init__.py
View file @ 4f1a5e52
......@@ -2,12 +2,28 @@ from mmdet.ops import (RoIAlign, SigmoidFocalLoss, get_compiler_version,
get_compiling_cuda_version, nms, roi_align,
sigmoid_focal_loss)
from .norm import NaiveSyncBatchNorm1d, NaiveSyncBatchNorm2d
from .sparse_block import (SparseBasicBlock, SparseBasicBlockV0,
SparseBottleneck, SparseBottleneckV0)
from .voxel import DynamicScatter, Voxelization, dynamic_scatter, voxelization
__all__ = [
'nms', 'soft_nms', 'RoIAlign', 'roi_align', 'get_compiler_version',
'get_compiling_cuda_version', 'build_conv_layer', 'NaiveSyncBatchNorm1d',
'NaiveSyncBatchNorm2d', 'batched_nms', 'Voxelization', 'voxelization',
'dynamic_scatter', 'DynamicScatter', 'sigmoid_focal_loss',
'SigmoidFocalLoss'
'nms',
'soft_nms',
'RoIAlign',
'roi_align',
'get_compiler_version',
'get_compiling_cuda_version',
'NaiveSyncBatchNorm1d',
'NaiveSyncBatchNorm2d',
'batched_nms',
'Voxelization',
'voxelization',
'dynamic_scatter',
'DynamicScatter',
'sigmoid_focal_loss',
'SigmoidFocalLoss',
'SparseBasicBlockV0',
'SparseBottleneckV0',
'SparseBasicBlock',
'SparseBottleneck',
]
mmdet3d/ops/iou3d/src/iou3d.cpp
View file @ 4f1a5e52
#include <torch/serialize/tensor.h>
#include <torch/extension.h>
#include <vector>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <torch/extension.h>
#include <torch/serialize/tensor.h>
#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))
#include <vector>
#define CHECK_ERROR(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
if (code != cudaSuccess)
{
fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
if (abort) exit(code);
}
#define CHECK_CUDA(x) \
TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDAtensor ")
#define CHECK_CONTIGUOUS(x) \
TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x)
#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))
#define CHECK_ERROR(ans) \
{ gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line,
bool abort = true) {
if (code != cudaSuccess) {
fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
line);
if (abort) exit(code);
}
}
const int THREADS_PER_BLOCK_NMS = sizeof(unsigned long long) * 8;
void boxesoverlapLauncher(const int num_a, const float *boxes_a, const int num_b, const float *boxes_b, float *ans_overlap);
void boxesioubevLauncher(const int num_a, const float *boxes_a, const int num_b, const float *boxes_b, float *ans_iou);
void nmsLauncher(const float *boxes, unsigned long long * mask, int boxes_num, float nms_overlap_thresh);
void nmsNormalLauncher(const float *boxes, unsigned long long * mask, int boxes_num, float nms_overlap_thresh);
int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_overlap){
// params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
// params boxes_b: (M, 5)
// params ans_overlap: (N, M)
CHECK_INPUT(boxes_a);
CHECK_INPUT(boxes_b);
CHECK_INPUT(ans_overlap);
int num_a = boxes_a.size(0);
int num_b = boxes_b.size(0);
const float * boxes_a_data = boxes_a.data<float>();
const float * boxes_b_data = boxes_b.data<float>();
float * ans_overlap_data = ans_overlap.data<float>();
boxesoverlapLauncher(num_a, boxes_a_data, num_b, boxes_b_data, ans_overlap_data);
return 1;
void boxesoverlapLauncher(const int num_a, const float *boxes_a,
const int num_b, const float *boxes_b,
float *ans_overlap);
void boxesioubevLauncher(const int num_a, const float *boxes_a, const int num_b,
const float *boxes_b, float *ans_iou);
void nmsLauncher(const float *boxes, unsigned long long *mask, int boxes_num,
float nms_overlap_thresh);
void nmsNormalLauncher(const float *boxes, unsigned long long *mask,
int boxes_num, float nms_overlap_thresh);
int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b,
at::Tensor ans_overlap) {
// params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
// params boxes_b: (M, 5)
// params ans_overlap: (N, M)
CHECK_INPUT(boxes_a);
CHECK_INPUT(boxes_b);
CHECK_INPUT(ans_overlap);
int num_a = boxes_a.size(0);
int num_b = boxes_b.size(0);
const float *boxes_a_data = boxes_a.data_ptr<float>();
const float *boxes_b_data = boxes_b.data_ptr<float>();
float *ans_overlap_data = ans_overlap.data_ptr<float>();
boxesoverlapLauncher(num_a, boxes_a_data, num_b, boxes_b_data,
ans_overlap_data);
return 1;
}
int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou){
// params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
// params boxes_b: (M, 5)
// params ans_overlap: (N, M)
int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b,
at::Tensor ans_iou) {
// params boxes_a: (N, 5) [x1, y1, x2, y2, ry]
// params boxes_b: (M, 5)
// params ans_overlap: (N, M)
CHECK_INPUT(boxes_a);
CHECK_INPUT(boxes_b);
CHECK_INPUT(ans_iou);
CHECK_INPUT(boxes_a);
CHECK_INPUT(boxes_b);
CHECK_INPUT(ans_iou);
int num_a = boxes_a.size(0);
int num_b = boxes_b.size(0);
int num_a = boxes_a.size(0);
int num_b = boxes_b.size(0);
const float * boxes_a_data = boxes_a.data<float>();
const float * boxes_b_data = boxes_b.data<float>();
float * ans_iou_data = ans_iou.data<float>();
const float *boxes_a_data = boxes_a.data_ptr<float>();
const float *boxes_b_data = boxes_b.data_ptr<float>();
float *ans_iou_data = ans_iou.data_ptr<float>();
boxesioubevLauncher(num_a, boxes_a_data, num_b, boxes_b_data, ans_iou_data);
boxesioubevLauncher(num_a, boxes_a_data, num_b, boxes_b_data, ans_iou_data);
return 1;
return 1;
}
int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
// params boxes: (N, 5) [x1, y1, x2, y2, ry]
// params keep: (N)
int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh) {
// params boxes: (N, 5) [x1, y1, x2, y2, ry]
// params keep: (N)
CHECK_INPUT(boxes);
CHECK_CONTIGUOUS(keep);
CHECK_INPUT(boxes);
CHECK_CONTIGUOUS(keep);
int boxes_num = boxes.size(0);
const float * boxes_data = boxes.data<float>();
long * keep_data = keep.data<long>();
int boxes_num = boxes.size(0);
const float *boxes_data = boxes.data_ptr<float>();
long *keep_data = keep.data_ptr<long>();
const int col_blocks = DIVUP(boxes_num, THREADS_PER_BLOCK_NMS);
const int col_blocks = DIVUP(boxes_num, THREADS_PER_BLOCK_NMS);
unsigned long long *mask_data = NULL;
CHECK_ERROR(cudaMalloc((void**)&mask_data, boxes_num * col_blocks * sizeof(unsigned long long)));
nmsLauncher(boxes_data, mask_data, boxes_num, nms_overlap_thresh);
unsigned long long *mask_data = NULL;
CHECK_ERROR(cudaMalloc((void **)&mask_data,
boxes_num * col_blocks * sizeof(unsigned long long)));
nmsLauncher(boxes_data, mask_data, boxes_num, nms_overlap_thresh);
// unsigned long long mask_cpu[boxes_num * col_blocks];
// unsigned long long *mask_cpu = new unsigned long long [boxes_num * col_blocks];
std::vector<unsigned long long> mask_cpu(boxes_num * col_blocks);
// unsigned long long mask_cpu[boxes_num * col_blocks];
// unsigned long long *mask_cpu = new unsigned long long [boxes_num *
// col_blocks];
std::vector<unsigned long long> mask_cpu(boxes_num * col_blocks);
// printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data, boxes_num * col_blocks * sizeof(unsigned long long),
cudaMemcpyDeviceToHost));
// printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data,
boxes_num * col_blocks * sizeof(unsigned long long),
cudaMemcpyDeviceToHost));
cudaFree(mask_data);
cudaFree(mask_data);
unsigned long long remv_cpu[col_blocks];
memset(remv_cpu, 0, col_blocks * sizeof(unsigned long long));
unsigned long long remv_cpu[col_blocks];
memset(remv_cpu, 0, col_blocks * sizeof(unsigned long long));
int num_to_keep = 0;
int num_to_keep = 0;
for (int i = 0; i < boxes_num; i++){
int nblock = i / THREADS_PER_BLOCK_NMS;
int inblock = i % THREADS_PER_BLOCK_NMS;
for (int i = 0; i < boxes_num; i++) {
int nblock = i / THREADS_PER_BLOCK_NMS;
int inblock = i % THREADS_PER_BLOCK_NMS;
if (!(remv_cpu[nblock] & (1ULL << inblock))){
keep_data[num_to_keep++] = i;
unsigned long long *p = &mask_cpu[0] + i * col_blocks;
for (int j = nblock; j < col_blocks; j++){
remv_cpu[j] |= p[j];
}
}
if (!(remv_cpu[nblock] & (1ULL << inblock))) {
keep_data[num_to_keep++] = i;
unsigned long long *p = &mask_cpu[0] + i * col_blocks;
for (int j = nblock; j < col_blocks; j++) {
remv_cpu[j] |= p[j];
}
}
if ( cudaSuccess != cudaGetLastError() ) printf( "Error!\n" );
}
if (cudaSuccess != cudaGetLastError()) printf("Error!\n");
return num_to_keep;
return num_to_keep;
}
int nms_normal_gpu(at::Tensor boxes, at::Tensor keep,
float nms_overlap_thresh) {
// params boxes: (N, 5) [x1, y1, x2, y2, ry]
// params keep: (N)
int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
// params boxes: (N, 5) [x1, y1, x2, y2, ry]
// params keep: (N)
CHECK_INPUT(boxes);
CHECK_CONTIGUOUS(keep);
CHECK_INPUT(boxes);
CHECK_CONTIGUOUS(keep);
int boxes_num = boxes.size(0);
const float * boxes_data = boxes.data<float>();
long * keep_data = keep.data<long>();
int boxes_num = boxes.size(0);
const float *boxes_data = boxes.data_ptr<float>();
long *keep_data = keep.data_ptr<long>();
const int col_blocks = DIVUP(boxes_num, THREADS_PER_BLOCK_NMS);
const int col_blocks = DIVUP(boxes_num, THREADS_PER_BLOCK_NMS);
unsigned long long *mask_data = NULL;
CHECK_ERROR(cudaMalloc((void**)&mask_data, boxes_num * col_blocks * sizeof(unsigned long long)));
nmsNormalLauncher(boxes_data, mask_data, boxes_num, nms_overlap_thresh);
unsigned long long *mask_data = NULL;
CHECK_ERROR(cudaMalloc((void **)&mask_data,
boxes_num * col_blocks * sizeof(unsigned long long)));
nmsNormalLauncher(boxes_data, mask_data, boxes_num, nms_overlap_thresh);
// unsigned long long mask_cpu[boxes_num * col_blocks];
// unsigned long long *mask_cpu = new unsigned long long [boxes_num * col_blocks];
std::vector<unsigned long long> mask_cpu(boxes_num * col_blocks);
// unsigned long long mask_cpu[boxes_num * col_blocks];
// unsigned long long *mask_cpu = new unsigned long long [boxes_num *
// col_blocks];
std::vector<unsigned long long> mask_cpu(boxes_num * col_blocks);
// printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data, boxes_num * col_blocks * sizeof(unsigned long long),
cudaMemcpyDeviceToHost));
// printf("boxes_num=%d, col_blocks=%d\n", boxes_num, col_blocks);
CHECK_ERROR(cudaMemcpy(&mask_cpu[0], mask_data,
boxes_num * col_blocks * sizeof(unsigned long long),
cudaMemcpyDeviceToHost));
cudaFree(mask_data);
cudaFree(mask_data);
unsigned long long remv_cpu[col_blocks];
memset(remv_cpu, 0, col_blocks * sizeof(unsigned long long));
unsigned long long remv_cpu[col_blocks];
memset(remv_cpu, 0, col_blocks * sizeof(unsigned long long));
int num_to_keep = 0;
int num_to_keep = 0;
for (int i = 0; i < boxes_num; i++){
int nblock = i / THREADS_PER_BLOCK_NMS;
int inblock = i % THREADS_PER_BLOCK_NMS;
for (int i = 0; i < boxes_num; i++) {
int nblock = i / THREADS_PER_BLOCK_NMS;
int inblock = i % THREADS_PER_BLOCK_NMS;
if (!(remv_cpu[nblock] & (1ULL << inblock))){
keep_data[num_to_keep++] = i;
unsigned long long *p = &mask_cpu[0] + i * col_blocks;
for (int j = nblock; j < col_blocks; j++){
remv_cpu[j] |= p[j];
}
}
if (!(remv_cpu[nblock] & (1ULL << inblock))) {
keep_data[num_to_keep++] = i;
unsigned long long *p = &mask_cpu[0] + i * col_blocks;
for (int j = nblock; j < col_blocks; j++) {
remv_cpu[j] |= p[j];
}
}
if ( cudaSuccess != cudaGetLastError() ) printf( "Error!\n" );
}
if (cudaSuccess != cudaGetLastError()) printf("Error!\n");
return num_to_keep;
return num_to_keep;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu, "oriented boxes overlap");
m.def("boxes_overlap_bev_gpu", &boxes_overlap_bev_gpu,
"oriented boxes overlap");
m.def("boxes_iou_bev_gpu", &boxes_iou_bev_gpu, "oriented boxes iou");
m.def("nms_gpu", &nms_gpu, "oriented nms gpu");
m.def("nms_normal_gpu", &nms_normal_gpu, "nms gpu");
......
mmdet3d/ops/norm.py
View file @ 4f1a5e52
import torch
import torch.distributed as dist
import torch.nn as nn
from mmcv.cnn import NORM_LAYERS
from torch.autograd.function import Function
from mmdet.ops.norm import norm_cfg
class AllReduce(Function):
......@@ -24,6 +23,7 @@ class AllReduce(Function):
return grad_output
@NORM_LAYERS.register_module('naiveSyncBN1d')
class NaiveSyncBatchNorm1d(nn.BatchNorm1d):
"""Syncronized Batch Normalization for 3D Tensors
......@@ -68,6 +68,7 @@ class NaiveSyncBatchNorm1d(nn.BatchNorm1d):
return input * scale + bias
@NORM_LAYERS.register_module('naiveSyncBN2d')
class NaiveSyncBatchNorm2d(nn.BatchNorm2d):
"""Syncronized Batch Normalization for 4D Tensors
......@@ -110,10 +111,3 @@ class NaiveSyncBatchNorm2d(nn.BatchNorm2d):
scale = scale.reshape(1, -1, 1, 1)
bias = bias.reshape(1, -1, 1, 1)
return input * scale + bias
norm_cfg.update({
'BN1d': ('bn', nn.BatchNorm1d),
'naiveSyncBN2d': ('bn', NaiveSyncBatchNorm2d),
'naiveSyncBN1d': ('bn', NaiveSyncBatchNorm1d),
})
mmdet3d/ops/roiaware_pool3d/points_in_boxes.py
View file @ 4f1a5e52
......@@ -4,12 +4,14 @@ from . import roiaware_pool3d_ext
def points_in_boxes_gpu(points, boxes):
"""
"""Find points that are in boxes (CUDA)
Args:
points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR coordinate
boxes (torch.Tensor): [B, T, 7],
num_valid_boxes <= T, [x, y, z, w, l, h, ry] in LiDAR coordinate,
(x, y, z) is the bottom center
Returns:
box_idxs_of_pts (torch.Tensor): (B, M), default background = -1
"""
......@@ -27,14 +29,20 @@ def points_in_boxes_gpu(points, boxes):
def points_in_boxes_cpu(points, boxes):
"""
"""Find points that are in boxes (CPU)
Note: Currently, the output of this function is different from that of
points_in_boxes_gpu.
Args:
points (torch.Tensor): [npoints, 3]
boxes (torch.Tensor): [N, 7], in LiDAR coordinate,
(x, y, z) is the bottom center
Returns:
point_indices (torch.Tensor): (N, npoints)
"""
# TODO: Refactor this function as a CPU version of points_in_boxes_gpu
assert boxes.shape[1] == 7
assert points.shape[1] == 3
......
mmdet3d/ops/roiaware_pool3d/roiaware_pool3d.py
View file @ 4f1a5e52
......@@ -10,7 +10,8 @@ class RoIAwarePool3d(nn.Module):
def __init__(self, out_size, max_pts_per_voxel=128, mode='max'):
super().__init__()
"""
"""RoIAwarePool3d module
Args:
out_size (int or tuple): n or [n1, n2, n3]
max_pts_per_voxel (int): m
......@@ -23,12 +24,14 @@ class RoIAwarePool3d(nn.Module):
self.mode = pool_method_map[mode]
def forward(self, rois, pts, pts_feature):
"""
"""RoIAwarePool3d module forward
Args:
rois (torch.Tensor): [N, 7],in LiDAR coordinate,
(x, y, z) is the bottom center of rois
pts (torch.Tensor): [npoints, 3]
pts_feature (torch.Tensor): [npoints, C]
Returns:
pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C]
"""
......@@ -43,7 +46,8 @@ class RoIAwarePool3dFunction(Function):
@staticmethod
def forward(ctx, rois, pts, pts_feature, out_size, max_pts_per_voxel,
mode):
"""
"""RoIAwarePool3d function forward
Args:
rois (torch.Tensor): [N, 7], in LiDAR coordinate,
(x, y, z) is the bottom center of rois
......@@ -52,6 +56,7 @@ class RoIAwarePool3dFunction(Function):
out_size (int or tuple): n or [n1, n2, n3]
max_pts_per_voxel (int): m
mode (int): 0 (max pool) or 1 (average pool)
Returns:
pooled_features (torch.Tensor): [N, out_x, out_y, out_z, C]
"""
......@@ -84,11 +89,12 @@ class RoIAwarePool3dFunction(Function):
@staticmethod
def backward(ctx, grad_out):
"""
"""RoIAwarePool3d function forward
Args:
grad_out: [N, out_x, out_y, out_z, C]
grad_out (torch.Tensor): [N, out_x, out_y, out_z, C]
Returns:
grad_in: [npoints, C]
grad_in (torch.Tensor): [npoints, C]
"""
ret = ctx.roiaware_pool3d_for_backward
pts_idx_of_voxels, argmax, mode, num_pts, num_channels = ret
......
mmdet3d/ops/roiaware_pool3d/src/points_in_boxes_cpu.cpp
View file @ 4f1a5e52
//Modified from
//https://github.com/sshaoshuai/PCDet/blob/master/pcdet/ops/roiaware_pool3d/src/roiaware_pool3d_kernel.cu
//Points in boxes cpu
//Written by Shaoshuai Shi
//All Rights Reserved 2019.
// Modified from
// https://github.com/sshaoshuai/PCDet/blob/master/pcdet/ops/roiaware_pool3d/src/roiaware_pool3d_kernel.cu
// Points in boxes cpu
// Written by Shaoshuai Shi
// All Rights Reserved 2019.
#include <torch/serialize/tensor.h>
#include <torch/extension.h>
#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <torch/extension.h>
#include <torch/serialize/tensor.h>
#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_CONTIGUOUS(x) \
TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
// #define DEBUG
inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz, float &local_x, float &local_y){
// should rotate pi/2 + alpha to translate LiDAR to local
float rot_angle = rz + M_PI / 2;
float cosa = cos(rot_angle), sina = sin(rot_angle);
local_x = shift_x * cosa + shift_y * (-sina);
local_y = shift_x * sina + shift_y * cosa;
inline void lidar_to_local_coords_cpu(float shift_x, float shift_y, float rz,
float &local_x, float &local_y) {
// should rotate pi/2 + alpha to translate LiDAR to local
float rot_angle = rz + M_PI / 2;
float cosa = cos(rot_angle), sina = sin(rot_angle);
local_x = shift_x * cosa + shift_y * (-sina);
local_y = shift_x * sina + shift_y * cosa;
}
inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d, float &local_x, float &local_y){
// param pt: (x, y, z)
// param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the bottom center
float x = pt[0], y = pt[1], z = pt[2];
float cx = box3d[0], cy = box3d[1], cz = box3d[2];
float w = box3d[3], l = box3d[4], h = box3d[5], rz = box3d[6];
cz += h / 2.0; // shift to the center since cz in box3d is the bottom center
if (fabsf(z - cz) > h / 2.0) return 0;
lidar_to_local_coords_cpu(x - cx, y - cy, rz, local_x, local_y);
float in_flag = (local_x > -l / 2.0) & (local_x < l / 2.0) & (local_y > -w / 2.0) & (local_y < w / 2.0);
return in_flag;
inline int check_pt_in_box3d_cpu(const float *pt, const float *box3d,
float &local_x, float &local_y) {
// param pt: (x, y, z)
// param box3d: (cx, cy, cz, w, l, h, rz) in LiDAR coordinate, cz in the
// bottom center
float x = pt[0], y = pt[1], z = pt[2];
float cx = box3d[0], cy = box3d[1], cz = box3d[2];
float w = box3d[3], l = box3d[4], h = box3d[5], rz = box3d[6];
cz += h / 2.0; // shift to the center since cz in box3d is the bottom center
if (fabsf(z - cz) > h / 2.0) return 0;
lidar_to_local_coords_cpu(x - cx, y - cy, rz, local_x, local_y);
float in_flag = (local_x > -l / 2.0) & (local_x < l / 2.0) &
(local_y > -w / 2.0) & (local_y < w / 2.0);
return in_flag;
}
int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor, at::Tensor pts_indices_tensor){
// params boxes: (N, 7) [x, y, z, w, l, h, rz] in LiDAR coordinate, z is the bottom center, each box DO NOT overlaps
// params pts: (npoints, 3) [x, y, z] in LiDAR coordinate
// params pts_indices: (N, npoints)
CHECK_CONTIGUOUS(boxes_tensor);
CHECK_CONTIGUOUS(pts_tensor);
CHECK_CONTIGUOUS(pts_indices_tensor);
int boxes_num = boxes_tensor.size(0);
int pts_num = pts_tensor.size(0);
const float *boxes = boxes_tensor.data<float>();
const float *pts = pts_tensor.data<float>();
int *pts_indices = pts_indices_tensor.data<int>();
float local_x = 0, local_y = 0;
for (int i = 0; i < boxes_num; i++){
for (int j = 0; j < pts_num; j++){
int cur_in_flag = check_pt_in_box3d_cpu(pts + j * 3, boxes + i * 7, local_x, local_y);
pts_indices[i * pts_num + j] = cur_in_flag;
}
int points_in_boxes_cpu(at::Tensor boxes_tensor, at::Tensor pts_tensor,
at::Tensor pts_indices_tensor) {
// params boxes: (N, 7) [x, y, z, w, l, h, rz] in LiDAR coordinate, z is the
// bottom center, each box DO NOT overlaps params pts: (npoints, 3) [x, y, z]
// in LiDAR coordinate params pts_indices: (N, npoints)
CHECK_CONTIGUOUS(boxes_tensor);
CHECK_CONTIGUOUS(pts_tensor);
CHECK_CONTIGUOUS(pts_indices_tensor);
int boxes_num = boxes_tensor.size(0);
int pts_num = pts_tensor.size(0);
const float *boxes = boxes_tensor.data_ptr<float>();
const float *pts = pts_tensor.data_ptr<float>();
int *pts_indices = pts_indices_tensor.data_ptr<int>();
float local_x = 0, local_y = 0;
for (int i = 0; i < boxes_num; i++) {
for (int j = 0; j < pts_num; j++) {
int cur_in_flag =
check_pt_in_box3d_cpu(pts + j * 3, boxes + i * 7, local_x, local_y);
pts_indices[i * pts_num + j] = cur_in_flag;
}
}
return 1;
return 1;
}
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