Commit aec41c7f authored by zhangwenwei's avatar zhangwenwei
Browse files

Merge master

parents 49f06039 4eca6606
from mmdet.datasets.registry import DATASETS
from mmdet.datasets.builder import DATASETS
from .builder import build_dataset
from .dataset_wrappers import RepeatFactorDataset
from .kitti2d_dataset import Kitti2DDataset
from .kitti_dataset import KittiDataset
from .loader import DistributedGroupSampler, GroupSampler, build_dataloader
from .nuscenes2d_dataset import NuScenes2DDataset
from .nuscenes_dataset import NuScenesDataset
from .pipelines import (GlobalRotScale, ObjectNoise, ObjectRangeFilter,
ObjectSample, PointShuffle, PointsRangeFilter,
RandomFlip3D)
__all__ = [
'KittiDataset', 'GroupSampler', 'DistributedGroupSampler',
'build_dataloader', 'RepeatFactorDataset', 'DATASETS', 'build_dataset',
'CocoDataset', 'Kitti2DDataset', 'NuScenesDataset', 'NuScenes2DDataset'
'CocoDataset', 'Kitti2DDataset', 'NuScenesDataset', 'ObjectSample',
'RandomFlip3D', 'ObjectNoise', 'GlobalRotScale', 'PointShuffle',
'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D'
]
import copy
from mmcv.utils import build_from_cfg
from mmdet.datasets import DATASETS, ConcatDataset, RepeatDataset
from mmdet.utils import build_from_cfg
from .dataset_wrappers import RepeatFactorDataset
......
......@@ -9,8 +9,8 @@ import torch
import torch.utils.data as torch_data
from mmdet.datasets import DATASETS
from mmdet.datasets.pipelines import Compose
from ..core.bbox import box_np_ops
from .pipelines import Compose
from .utils import remove_dontcare
......
from pycocotools.coco import COCO
from mmdet3d.core.evaluation.coco_utils import getImgIds
from mmdet.datasets import DATASETS, CocoDataset
@DATASETS.register_module
class NuScenes2DDataset(CocoDataset):
CLASSES = ('car', 'truck', 'trailer', 'bus', 'construction_vehicle',
'bicycle', 'motorcycle', 'pedestrian', 'traffic_cone',
'barrier')
def load_annotations(self, ann_file):
if not self.class_names:
self.class_names = self.CLASSES
self.coco = COCO(ann_file)
# send class_names into the get id
# in case we only need to train on several classes
# by default self.class_names = CLASSES
self.cat_ids = self.coco.getCatIds(catNms=self.class_names)
self.cat2label = {
cat_id: i # + 1 rm +1 here thus the 0-79 are fg, 80 is bg
for i, cat_id in enumerate(self.cat_ids)
}
# send cat ids to the get img id
# in case we only need to train on several classes
if len(self.cat_ids) < len(self.CLASSES):
self.img_ids = getImgIds(self.coco, catIds=self.cat_ids)
else:
self.img_ids = self.coco.getImgIds()
img_infos = []
for i in self.img_ids:
info = self.coco.loadImgs([i])[0]
info['filename'] = info['file_name']
img_infos.append(info)
return img_infos
......@@ -9,8 +9,8 @@ import torch.utils.data as torch_data
from nuscenes.utils.data_classes import Box as NuScenesBox
from mmdet.datasets import DATASETS
from mmdet.datasets.pipelines import Compose
from ..core.bbox import box_np_ops
from .pipelines import Compose
@DATASETS.register_module
......
from mmdet.datasets.pipelines import Compose
from .dbsampler import DataBaseSampler, MMDataBaseSampler
from .formating import DefaultFormatBundle, DefaultFormatBundle3D
from .loading import LoadMultiViewImageFromFiles, LoadPointsFromFile
from .train_aug import (GlobalRotScale, ObjectNoise, ObjectRangeFilter,
ObjectSample, PointShuffle, PointsRangeFilter,
RandomFlip3D)
__all__ = [
'ObjectSample', 'RandomFlip3D', 'ObjectNoise', 'GlobalRotScale',
'PointShuffle', 'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D'
'PointShuffle', 'ObjectRangeFilter', 'PointsRangeFilter', 'Collect3D',
'Compose', 'LoadMultiViewImageFromFiles', 'LoadPointsFromFile',
'DefaultFormatBundle', 'DefaultFormatBundle3D', 'DataBaseSampler',
'MMDataBaseSampler'
]
......@@ -68,7 +68,7 @@ class DataBaseSampler(object):
db_infos = pickle.load(f)
# filter database infos
from mmdet3d.apis import get_root_logger
from mmdet.apis import get_root_logger
logger = get_root_logger()
for k, v in db_infos.items():
logger.info(f'load {len(v)} {k} database infos')
......
import numpy as np
from mmcv.parallel import DataContainer as DC
from mmdet.datasets.builder import PIPELINES
from mmdet.datasets.pipelines import to_tensor
from mmdet.datasets.registry import PIPELINES
PIPELINES._module_dict.pop('DefaultFormatBundle')
......
......@@ -3,7 +3,7 @@ import os.path as osp
import mmcv
import numpy as np
from mmdet.datasets.registry import PIPELINES
from mmdet.datasets.builder import PIPELINES
@PIPELINES.register_module
......
import numpy as np
from mmcv.utils import build_from_cfg
from mmdet3d.core.bbox import box_np_ops
from mmdet3d.utils import build_from_cfg
from mmdet.datasets.builder import PIPELINES
from mmdet.datasets.pipelines import RandomFlip
from mmdet.datasets.registry import PIPELINES
from ..registry import OBJECTSAMPLERS
from .data_augment_utils import noise_per_object_v3_
......
from mmdet.utils import Registry
from mmcv.utils import Registry
OBJECTSAMPLERS = Registry('Object sampler')
import numpy as np
import torch
from mmcv.cnn import normal_init
from mmcv.cnn import bias_init_with_prob, normal_init
from mmdet3d.core import box_torch_ops, boxes3d_to_bev_torch_lidar
from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
from mmdet.models import HEADS
from ..utils import bias_init_with_prob
from .second_head import SECONDHead
......@@ -15,12 +14,6 @@ class Anchor3DVeloHead(SECONDHead):
Args:
in_channels (int): Number of channels in the input feature map.
feat_channels (int): Number of channels of the feature map.
anchor_scales (Iterable): Anchor scales.
anchor_ratios (Iterable): Anchor aspect ratios.
anchor_strides (Iterable): Anchor strides.
anchor_base_sizes (Iterable): Anchor base sizes.
target_means (Iterable): Mean values of regression targets.
target_stds (Iterable): Std values of regression targets.
loss_cls (dict): Config of classification loss.
loss_bbox (dict): Config of localization loss.
""" # noqa: W605
......@@ -31,25 +24,25 @@ class Anchor3DVeloHead(SECONDHead):
in_channels,
train_cfg,
test_cfg,
cache_anchor=False,
feat_channels=256,
use_direction_classifier=True,
encode_bg_as_zeros=False,
box_code_size=9,
anchor_generator=dict(type='AnchorGeneratorRange', ),
anchor_range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
anchor_strides=[2],
anchor_sizes=[[1.6, 3.9, 1.56]],
anchor_rotations=[0, 1.57],
anchor_custom_values=[0, 0],
anchor_generator=dict(
type='Anchor3DRangeGenerator',
range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
strides=[2],
sizes=[[1.6, 3.9, 1.56]],
rotations=[0, 1.57],
custom_values=[0, 0],
reshape_out=True,
),
assigner_per_size=False,
assign_per_class=False,
diff_rad_by_sin=True,
dir_offset=0,
dir_limit_offset=1,
target_means=(.0, .0, .0, .0),
target_stds=(1.0, 1.0, 1.0, 1.0),
bbox_coder=dict(type='Residual3DBoxCoder', ),
bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
loss_cls=dict(
type='CrossEntropyLoss',
use_sigmoid=True,
......@@ -58,14 +51,11 @@ class Anchor3DVeloHead(SECONDHead):
type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=2.0),
loss_dir=dict(type='CrossEntropyLoss', loss_weight=0.2)):
super().__init__(class_names, in_channels, train_cfg, test_cfg,
cache_anchor, feat_channels, use_direction_classifier,
feat_channels, use_direction_classifier,
encode_bg_as_zeros, box_code_size, anchor_generator,
anchor_range, anchor_strides, anchor_sizes,
anchor_rotations, anchor_custom_values,
assigner_per_size, assign_per_class, diff_rad_by_sin,
dir_offset, dir_limit_offset, target_means,
target_stds, bbox_coder, loss_cls, loss_bbox,
loss_dir)
dir_offset, dir_limit_offset, bbox_coder, loss_cls,
loss_bbox, loss_dir)
self.num_classes = num_classes
# build head layers & losses
if not self.use_sigmoid_cls:
......@@ -131,9 +121,7 @@ class Anchor3DVeloHead(SECONDHead):
scores = scores[topk_inds, :]
dir_cls_score = dir_cls_score[topk_inds]
bboxes = self.bbox_coder.decode_torch(anchors, bbox_pred,
self.target_means,
self.target_stds)
bboxes = self.bbox_coder.decode(anchors, bbox_pred)
mlvl_bboxes.append(bboxes)
mlvl_scores.append(scores)
mlvl_dir_scores.append(dir_cls_score)
......
from __future__ import division
import numpy as np
import torch
import torch.nn as nn
from mmcv.cnn import normal_init
from mmcv.cnn import bias_init_with_prob, normal_init
from mmdet3d.core import (PseudoSampler, box_torch_ops,
boxes3d_to_bev_torch_lidar, build_anchor_generator,
......@@ -12,24 +10,37 @@ from mmdet3d.core import (PseudoSampler, box_torch_ops,
from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu
from mmdet.models import HEADS
from ..builder import build_loss
from ..utils import bias_init_with_prob
from .train_mixins import AnchorTrainMixin
@HEADS.register_module
class SECONDHead(nn.Module, AnchorTrainMixin):
"""Anchor-based head (RPN, RetinaNet, SSD, etc.).
"""Anchor-based head for VoxelNet detectors.
Args:
class_name (list[str]): name of classes (TODO: to be removed)
in_channels (int): Number of channels in the input feature map.
train_cfg (dict): train configs
test_cfg (dict): test configs
feat_channels (int): Number of channels of the feature map.
anchor_scales (Iterable): Anchor scales.
anchor_ratios (Iterable): Anchor aspect ratios.
anchor_strides (Iterable): Anchor strides.
anchor_base_sizes (Iterable): Anchor base sizes.
target_means (Iterable): Mean values of regression targets.
target_stds (Iterable): Std values of regression targets.
use_direction_classifier (bool): Whether to add a direction classifier.
encode_bg_as_zeros (bool): Whether to use sigmoid of softmax
(TODO: to be removed)
box_code_size (int): The size of box code.
anchor_generator(dict): Config dict of anchor generator.
assigner_per_size (bool): Whether to do assignment for each separate
anchor size.
assign_per_class (bool): Whether to do assignment for each class.
diff_rad_by_sin (bool): Whether to change the difference into sin
difference for box regression loss.
dir_offset (float | int): The offset of BEV rotation angles
(TODO: may be moved into box coder)
dirlimit_offset (float | int): The limited range of BEV rotation angles
(TODO: may be moved into box coder)
box_coder (dict): Config dict of box coders.
loss_cls (dict): Config of classification loss.
loss_bbox (dict): Config of localization loss.
loss_dir (dict): Config of direction classifier loss.
""" # noqa: W605
def __init__(self,
......@@ -37,25 +48,24 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
in_channels,
train_cfg,
test_cfg,
cache_anchor=False,
feat_channels=256,
use_direction_classifier=True,
encode_bg_as_zeros=False,
box_code_size=7,
anchor_generator=dict(type='AnchorGeneratorRange'),
anchor_range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
anchor_strides=[2],
anchor_sizes=[[1.6, 3.9, 1.56]],
anchor_rotations=[0, 1.57],
anchor_custom_values=[],
anchor_generator=dict(
type='Anchor3DRangeGenerator',
range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
strides=[2],
sizes=[[1.6, 3.9, 1.56]],
rotations=[0, 1.57],
custom_values=[],
reshape_out=False),
assigner_per_size=False,
assign_per_class=False,
diff_rad_by_sin=True,
dir_offset=0,
dir_limit_offset=1,
target_means=(.0, .0, .0, .0),
target_stds=(1.0, 1.0, 1.0, 1.0),
bbox_coder=dict(type='Residual3DBoxCoder'),
bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
loss_cls=dict(
type='CrossEntropyLoss',
use_sigmoid=True,
......@@ -94,29 +104,9 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
]
# build anchor generator
self.anchor_range = anchor_range
self.anchor_rotations = anchor_rotations
self.anchor_strides = anchor_strides
self.anchor_sizes = anchor_sizes
self.target_means = target_means
self.target_stds = target_stds
self.anchor_generators = []
self.anchor_generator = build_anchor_generator(anchor_generator)
# In 3D detection, the anchor stride is connected with anchor size
self.num_anchors = (
len(self.anchor_rotations) * len(self.anchor_sizes))
# if len(self.anchor_sizes) != self.anchor_strides:
# # this means different anchor in the same anchor strides
# anchor_sizes = [self.anchor_sizes]
for anchor_stride in self.anchor_strides:
anchor_generator.update(
anchor_ranges=anchor_range,
sizes=self.anchor_sizes,
stride=anchor_stride,
rotations=anchor_rotations,
custom_values=anchor_custom_values,
cache_anchor=cache_anchor)
self.anchor_generators.append(
build_anchor_generator(anchor_generator))
self.num_anchors = self.anchor_generator.num_base_anchors
self._init_layers()
self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
......@@ -152,7 +142,7 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
def forward(self, feats):
return multi_apply(self.forward_single, feats)
def get_anchors(self, featmap_sizes, input_metas):
def get_anchors(self, featmap_sizes, input_metas, device='cuda'):
"""Get anchors according to feature map sizes.
Args:
featmap_sizes (list[tuple]): Multi-level feature map sizes.
......@@ -161,16 +151,10 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
tuple: anchors of each image, valid flags of each image
"""
num_imgs = len(input_metas)
num_levels = len(featmap_sizes)
# since feature map sizes of all images are the same, we only compute
# anchors for one time
multi_level_anchors = []
for i in range(num_levels):
anchors = self.anchor_generators[i].grid_anchors(featmap_sizes[i])
if not self.assigner_per_size:
anchors = anchors.reshape(-1, anchors.size(-1))
multi_level_anchors.append(anchors)
multi_level_anchors = self.anchor_generator.grid_anchors(
featmap_sizes, device=device)
anchor_list = [multi_level_anchors for _ in range(num_imgs)]
return anchor_list
......@@ -237,16 +221,15 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
input_metas,
gt_bboxes_ignore=None):
featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
assert len(featmap_sizes) == len(self.anchor_generators)
anchor_list = self.get_anchors(featmap_sizes, input_metas)
assert len(featmap_sizes) == self.anchor_generator.num_levels
device = cls_scores[0].device
anchor_list = self.get_anchors(
featmap_sizes, input_metas, device=device)
label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1
cls_reg_targets = self.anchor_target_3d(
anchor_list,
gt_bboxes,
input_metas,
self.target_means,
self.target_stds,
gt_bboxes_ignore_list=gt_bboxes_ignore,
gt_labels_list=gt_labels,
num_classes=self.num_classes,
......@@ -288,12 +271,14 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
assert len(cls_scores) == len(bbox_preds)
assert len(cls_scores) == len(dir_cls_preds)
num_levels = len(cls_scores)
featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
device = cls_scores[0].device
mlvl_anchors = self.anchor_generator.grid_anchors(
featmap_sizes, device=device)
mlvl_anchors = [
self.anchor_generators[i].grid_anchors(
cls_scores[i].size()[-2:]).reshape(-1, self.box_code_size)
for i in range(num_levels)
anchor.reshape(-1, self.box_code_size) for anchor in mlvl_anchors
]
result_list = []
for img_id in range(len(input_metas)):
cls_score_list = [
......@@ -353,9 +338,7 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
bbox_pred = bbox_pred[thr_inds]
scores = scores[thr_inds]
dir_cls_scores = dir_cls_score[thr_inds]
bboxes = self.bbox_coder.decode_torch(anchors, bbox_pred,
self.target_means,
self.target_stds)
bboxes = self.bbox_coder.decode(anchors, bbox_pred)
bboxes_for_nms = boxes3d_to_bev_torch_lidar(bboxes)
mlvl_bboxes_for_nms.append(bboxes_for_nms)
mlvl_bboxes.append(bboxes)
......@@ -383,6 +366,7 @@ class SECONDHead(nn.Module, AnchorTrainMixin):
selected_scores = mlvl_scores[selected]
selected_label_preds = mlvl_label_preds[selected]
selected_dir_scores = mlvl_dir_scores[selected]
# TODO: move dir_offset to box coder
dir_rot = box_torch_ops.limit_period(
selected_bboxes[..., -1] - self.dir_offset,
self.dir_limit_offset, np.pi)
......
......@@ -11,8 +11,6 @@ class AnchorTrainMixin(object):
anchor_list,
gt_bboxes_list,
input_metas,
target_means,
target_stds,
gt_bboxes_ignore_list=None,
gt_labels_list=None,
label_channels=1,
......@@ -24,8 +22,6 @@ class AnchorTrainMixin(object):
anchor_list (list[list]): Multi level anchors of each image.
gt_bboxes_list (list[Tensor]): Ground truth bboxes of each image.
img_metas (list[dict]): Meta info of each image.
target_means (Iterable): Mean value of regression targets.
target_stds (Iterable): Std value of regression targets.
Returns:
tuple
......@@ -57,8 +53,6 @@ class AnchorTrainMixin(object):
gt_bboxes_ignore_list,
gt_labels_list,
input_metas,
target_means=target_means,
target_stds=target_stds,
label_channels=label_channels,
num_classes=num_classes,
sampling=sampling)
......@@ -89,8 +83,6 @@ class AnchorTrainMixin(object):
gt_bboxes_ignore,
gt_labels,
input_meta,
target_means,
target_stds,
label_channels=1,
num_classes=1,
sampling=True):
......@@ -111,13 +103,12 @@ class AnchorTrainMixin(object):
anchor_targets = self.anchor_target_single_assigner(
assigner, current_anchors, gt_bboxes[gt_per_cls, :],
gt_bboxes_ignore, gt_labels[gt_per_cls], input_meta,
target_means, target_stds, label_channels, num_classes,
sampling)
label_channels, num_classes, sampling)
else:
anchor_targets = self.anchor_target_single_assigner(
assigner, current_anchors, gt_bboxes, gt_bboxes_ignore,
gt_labels, input_meta, target_means, target_stds,
label_channels, num_classes, sampling)
gt_labels, input_meta, label_channels, num_classes,
sampling)
(labels, label_weights, bbox_targets, bbox_weights,
dir_targets, dir_weights, pos_inds, neg_inds) = anchor_targets
......@@ -156,8 +147,7 @@ class AnchorTrainMixin(object):
else:
return self.anchor_target_single_assigner(
self.bbox_assigner, anchors, gt_bboxes, gt_bboxes_ignore,
gt_labels, input_meta, target_means, target_stds,
label_channels, num_classes, sampling)
gt_labels, input_meta, label_channels, num_classes, sampling)
def anchor_target_single_assigner(self,
bbox_assigner,
......@@ -166,8 +156,6 @@ class AnchorTrainMixin(object):
gt_bboxes_ignore,
gt_labels,
input_meta,
target_means,
target_stds,
label_channels=1,
num_classes=1,
sampling=True):
......@@ -188,18 +176,17 @@ class AnchorTrainMixin(object):
neg_inds = sampling_result.neg_inds
else:
pos_inds = torch.nonzero(
anchors.new_zeros((anchors.shape[0], ), dtype=torch.long) > 0
anchors.new_zeros((anchors.shape[0], ), dtype=torch.bool) > 0
).squeeze(-1).unique()
neg_inds = torch.nonzero(
anchors.new_zeros((anchors.shape[0], ), dtype=torch.long) ==
anchors.new_zeros((anchors.shape[0], ), dtype=torch.bool) ==
0).squeeze(-1).unique()
if gt_labels is not None:
labels += num_classes
if len(pos_inds) > 0:
pos_bbox_targets = self.bbox_coder.encode_torch(
sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes,
target_means, target_stds)
pos_bbox_targets = self.bbox_coder.encode(
sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes)
pos_dir_targets = get_direction_target(
sampling_result.pos_bboxes,
pos_bbox_targets,
......
from mmdet.models.builder import build
from mmdet.models.registry import (BACKBONES, DETECTORS, HEADS, LOSSES, NECKS,
ROI_EXTRACTORS, SHARED_HEADS)
from mmdet.models.builder import (BACKBONES, DETECTORS, HEADS, LOSSES, NECKS,
ROI_EXTRACTORS, SHARED_HEADS, build)
from .registry import FUSION_LAYERS, MIDDLE_ENCODERS, VOXEL_ENCODERS
......
......@@ -235,9 +235,8 @@ class PointFusion(nn.Module):
pts.new_tensor(img_meta['pcd_trans'])
if 'pcd_trans' in img_meta.keys() else 0)
pcd_rotate_mat = (
pts.new_tensor(img_meta['pcd_rotation'])
if 'pcd_rotation' in img_meta.keys() else
torch.eye(3).type_as(pts).to(pts.device))
pts.new_tensor(img_meta['pcd_rotation']) if 'pcd_rotation'
in img_meta.keys() else torch.eye(3).type_as(pts).to(pts.device))
img_scale_factor = (
img_meta['scale_factor']
if 'scale_factor' in img_meta.keys() else 1)
......
from mmdet.utils import Registry
from mmcv.utils import Registry
VOXEL_ENCODERS = Registry('voxel_encoder')
MIDDLE_ENCODERS = Registry('middle_encoder')
......
from mmdet.models.utils import ResLayer, bias_init_with_prob
__all__ = ['bias_init_with_prob', 'ResLayer']
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
#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);
#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){
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)
......@@ -40,16 +51,18 @@ int boxes_overlap_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans
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>();
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);
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){
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)
......@@ -61,16 +74,16 @@ int boxes_iou_bev_gpu(at::Tensor boxes_a, at::Tensor boxes_b, at::Tensor ans_iou
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);
return 1;
}
int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
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)
......@@ -78,21 +91,24 @@ int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
CHECK_CONTIGUOUS(keep);
int boxes_num = boxes.size(0);
const float * boxes_data = boxes.data<float>();
long * keep_data = keep.data<long>();
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);
unsigned long long *mask_data = NULL;
CHECK_ERROR(cudaMalloc((void**)&mask_data, boxes_num * col_blocks * sizeof(unsigned long long)));
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];
// 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),
// 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);
......@@ -102,25 +118,25 @@ int nms_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
int num_to_keep = 0;
for (int i = 0; i < boxes_num; i++){
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))){
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++){
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;
}
int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
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)
......@@ -128,21 +144,24 @@ int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
CHECK_CONTIGUOUS(keep);
int boxes_num = boxes.size(0);
const float * boxes_data = boxes.data<float>();
long * keep_data = keep.data<long>();
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);
unsigned long long *mask_data = NULL;
CHECK_ERROR(cudaMalloc((void**)&mask_data, boxes_num * col_blocks * sizeof(unsigned long long)));
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];
// 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),
// 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);
......@@ -152,27 +171,26 @@ int nms_normal_gpu(at::Tensor boxes, at::Tensor keep, float nms_overlap_thresh){
int num_to_keep = 0;
for (int i = 0; i < boxes_num; i++){
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))){
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++){
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;
}
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");
......
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