@ -191,36 +191,34 @@ class YOLOV3Head(BaseDenseHead, BBoxTestMixin):
Returns :
list [ tuple [ Tensor , Tensor ] ] : Each item in result_list is 2 - tuple .
The first item is an ( n , 5 ) tensor , where the first 4 columns
are bounding box positions ( tl_x , tl_y , br_x , br_y ) and the
5 - th column is a score between 0 and 1. The second item is a
( n , ) tensor where each item is the predicted class label of the
corresponding box .
The first item is an ( n , 5 ) tensor , where 5 represent
( tl_x , tl_y , br_x , br_y , score ) and the score between 0 and 1.
The shape of the second tensor in the tuple is ( n , ) , and
each element represents the class label of the corresponding
box .
"""
result_list = [ ]
num_levels = len ( pred_maps )
for img_id in range ( len ( img_metas ) ) :
pred_maps_list = [
pred_maps [ i ] [ img_id ] . detach ( ) for i in range ( num_levels )
]
scale_factor = img_metas [ img_id ] [ ' scale_factor ' ]
proposals = self . _get_bboxes_single ( pred_maps_list , scale_factor ,
cfg , rescale , with_nms )
result_list . append ( proposals )
pred_maps_list = [ pred_maps [ i ] . detach ( ) for i in range ( num_levels ) ]
scale_factors = [
img_metas [ i ] [ ' scale_factor ' ]
for i in range ( pred_maps_list [ 0 ] . shape [ 0 ] )
]
result_list = self . _get_bboxes ( pred_maps_list , scale_factors , cfg ,
rescale , with_nms )
return result_list
def _get_bboxes_single ( self ,
pred_maps_list ,
scale_factor ,
cfg ,
rescale = False ,
with_nms = True ) :
def _get_bboxes ( self ,
pred_maps_list ,
scale_factors ,
cfg ,
rescale = False ,
with_nms = True ) :
""" Transform outputs for a single batch item into bbox predictions.
Args :
pred_maps_list ( list [ Tensor ] ) : Prediction maps for different scales
of each single image in the batch .
scale_factor ( ndarray ) : Scale factor of the image arrange as
scale_factors ( list ( ndarray ) ) : Scale factor of the image arrange as
( w_scale , h_scale , w_scale , h_scale ) .
cfg ( mmcv . Config | None ) : Test / postprocessing configuration ,
if None , test_cfg would be used .
@ -230,62 +228,71 @@ class YOLOV3Head(BaseDenseHead, BBoxTestMixin):
Default : True .
Returns :
tuple ( Tensor ) :
det_bboxes ( Tensor ) : BBox predictions in shape ( n , 5 ) , where
the first 4 columns are bounding box positions
( tl_x , tl_y , br_x , br_y ) and the 5 - th column is a score
between 0 and 1.
det_labels ( Tensor ) : A ( n , ) tensor where each item is the
predicted class label of the corresponding box .
list [ tuple [ Tensor , Tensor ] ] : Each item in result_list is 2 - tuple .
The first item is an ( n , 5 ) tensor , where 5 represent
( tl_x , tl_y , br_x , br_y , score ) and the score between 0 and 1.
The shape of the second tensor in the tuple is ( n , ) , and
each element represents the class label of the corresponding
box .
"""
cfg = self . test_cfg if cfg is None else cfg
assert len ( pred_maps_list ) == self . num_levels
multi_lvl_bboxes = [ ]
multi_lvl_cls_scores = [ ]
multi_lvl_conf_scores = [ ]
num_levels = len ( pred_maps_list )
device = pred_maps_list [ 0 ] . device
batch_size = pred_maps_list [ 0 ] . shape [ 0 ]
featmap_sizes = [
pred_maps_list [ i ] . shape [ - 2 : ] for i in range ( num_levels )
pred_maps_list [ i ] . shape [ - 2 : ] for i in range ( self . num_levels )
]
multi_lvl_anchors = self . anchor_generator . grid_anchors (
featmap_sizes , pred_maps_list [ 0 ] [ 0 ] . device )
featmap_sizes , device )
# convert to tensor to keep tracing
nms_pre_tensor = torch . tensor (
cfg . get ( ' nms_pre ' , - 1 ) , device = device , dtype = torch . long )
multi_lvl_bboxes = [ ]
multi_lvl_cls_scores = [ ]
multi_lvl_conf_scores = [ ]
for i in range ( self . num_levels ) :
# get some key info for current scale
pred_map = pred_maps_list [ i ]
stride = self . featmap_strides [ i ]
# (h, w, num_anchors*num_attrib) -> (h*w*num_anchors, num_attrib)
pred_map = pred_map . permute ( 1 , 2 , 0 ) . reshape ( - 1 , self . num_attrib )
pred_map [ . . . , : 2 ] = torch . sigmoid ( pred_map [ . . . , : 2 ] )
bbox_pred = self . bbox_coder . decode ( multi_lvl_anchors [ i ] ,
pred_map [ . . . , : 4 ] , stride )
# (b,h, w, num_anchors*num_attrib) ->
# (b,h*w*num_anchors, num_attrib)
pred_map = pred_map . permute ( 0 , 2 , 3 ,
1 ) . reshape ( batch_size , - 1 ,
self . num_attrib )
# Inplace operation like
# ```pred_map[..., :2] = \torch.sigmoid(pred_map[..., :2])```
# would create constant tensor when exporting to onnx
pred_map_conf = torch . sigmoid ( pred_map [ . . . , : 2 ] )
pred_map_rest = pred_map [ . . . , 2 : ]
pred_map = torch . cat ( [ pred_map_conf , pred_map_rest ] , dim = - 1 )
pred_map_boxes = pred_map [ . . . , : 4 ]
multi_lvl_anchor = multi_lvl_anchors [ i ]
multi_lvl_anchor = multi_lvl_anchor . expand_as ( pred_map_boxes )
bbox_pred = self . bbox_coder . decode ( multi_lvl_anchor ,
pred_map_boxes , stride )
# conf and cls
conf_pred = torch . sigmoid ( pred_map [ . . . , 4 ] ) . view ( - 1 )
conf_pred = torch . sigmoid ( pred_map [ . . . , 4 ] )
cls_pred = torch . sigmoid ( pred_map [ . . . , 5 : ] ) . view (
- 1 , self . num_classes ) # Cls pred one-hot.
# Filtering out all predictions with conf < conf_thr
conf_thr = cfg . get ( ' conf_thr ' , - 1 )
if conf_thr > 0 and ( not torch . onnx . is_in_onnx_export ( ) ) :
# TensorRT not support NonZero
# add as_tuple=False for compatibility in Pytorch 1.6
# flatten would create a Reshape op with constant values,
# and raise RuntimeError when doing inference in ONNX Runtime
# with a different input image (#4221).
conf_inds = conf_pred . ge ( conf_thr ) . nonzero (
as_tuple = False ) . squeeze ( 1 )
bbox_pred = bbox_pred [ conf_inds , : ]
cls_pred = cls_pred [ conf_inds , : ]
conf_pred = conf_pred [ conf_inds ]
batch_size , - 1 , self . num_classes ) # Cls pred one-hot.
# Get top-k prediction
nms_pre = cfg . get ( ' nms_pre ' , - 1 )
if 0 < nms_pre < conf_pred . size ( 0 ) :
# Always keep topk op for dynamic input in onnx
if nms_pre_tensor > 0 and ( torch . onnx . is_in_onnx_export ( )
or conf_pred . shape [ 1 ] > nms_pre_tensor ) :
from torch import _shape_as_tensor
# keep shape as tensor and get k
num_anchor = _shape_as_tensor ( conf_pred ) [ 1 ] . to ( device )
nms_pre = torch . where ( nms_pre_tensor < num_anchor ,
nms_pre_tensor , num_anchor )
_ , topk_inds = conf_pred . topk ( nms_pre )
bbox_pred = bbox_pred [ topk_inds , : ]
cls_pred = cls_pred [ topk_inds , : ]
conf_pred = conf_pred [ topk_inds ]
batch_inds = torch . arange ( batch_size ) . view (
- 1 , 1 ) . expand_as ( topk_inds ) . long ( )
bbox_pred = bbox_pred [ batch_inds , topk_inds , : ]
cls_pred = cls_pred [ batch_inds , topk_inds , : ]
conf_pred = conf_pred [ batch_inds , topk_inds ]
# Save the result of current scale
multi_lvl_bboxes . append ( bbox_pred )
@ -293,43 +300,70 @@ class YOLOV3Head(BaseDenseHead, BBoxTestMixin):
multi_lvl_conf_scores . append ( conf_pred )
# Merge the results of different scales together
multi_lvl_bboxes = torch . cat ( multi_lvl_bboxes )
multi_lvl_cls_scores = torch . cat ( multi_lvl_cls_scores )
multi_lvl_conf_scores = torch . cat ( multi_lvl_conf_scores )
batch_mlvl_bboxes = torch . cat ( multi_lvl_bboxes , dim = 1 )
batch_mlvl_scores = torch . cat ( multi_lvl_cls_scores , dim = 1 )
batch_mlvl_conf_scores = torch . cat ( multi_lvl_conf_scores , dim = 1 )
# Set max number of box to be feed into nms in deployment
deploy_nms_pre = cfg . get ( ' deploy_nms_pre ' , - 1 )
if deploy_nms_pre > 0 and torch . onnx . is_in_onnx_export ( ) :
_ , topk_inds = multi_lvl_conf_scores . topk ( deploy_nms_pre )
multi_lvl_bboxes = multi_lvl_bboxes [ topk_inds , : ]
multi_lvl_cls_scores = multi_lvl_cls_scores [ topk_inds , : ]
multi_lvl_conf_scores = multi_lvl_conf_scores [ topk_inds ]
if with_nms and ( multi_lvl_conf_scores . size ( 0 ) == 0 ) :
_ , topk_inds = batch_mlvl_conf_scores . topk ( deploy_nms_pre )
batch_inds = torch . arange ( batch_size ) . view (
- 1 , 1 ) . expand_as ( topk_inds ) . long ( )
batch_mlvl_bboxes = batch_mlvl_bboxes [ batch_inds , topk_inds , : ]
batch_mlvl_scores = batch_mlvl_scores [ batch_inds , topk_inds , : ]
batch_mlvl_conf_scores = batch_mlvl_conf_scores [ batch_inds ,
topk_inds ]
if with_nms and ( batch_mlvl_conf_scores . size ( 0 ) == 0 ) :
return torch . zeros ( ( 0 , 5 ) ) , torch . zeros ( ( 0 , ) )
if rescale :
multi_lvl_bboxes / = multi_lvl_bboxes . new_tensor ( scale_factor )
batch_mlvl_bboxes / = batch_mlvl_bboxes . new_tensor (
scale_factors ) . unsqueeze ( 1 )
# In mmdet 2.x, the class_id for background is num_classes.
# i.e., the last column.
padding = multi_lvl_cls_scores . new_zeros ( multi_lvl_cls_scores . shape [ 0 ] ,
1 )
multi_lvl_cls_scores = torch . cat ( [ multi_lvl_cls_scores , padding ] ,
dim = 1 )
padding = batch_mlvl_scores . new_zeros ( batch_size ,
batch_mlvl_scores . shape [ 1 ] , 1 )
batch_mlvl_scores = torch . cat ( [ batch_mlvl_scores , padding ] , dim = - 1 )
# Support exporting to onnx without nms
if with_nms and cfg . get ( ' nms ' , None ) is not None :
det_bboxes , det_labels = multiclass_nms (
multi_lvl_bboxes ,
multi_lvl_cls_scores ,
cfg . score_thr ,
cfg . nms ,
cfg . max_per_img ,
score_factors = multi_lvl_conf_scores )
return det_bboxes , det_labels
det_results = [ ]
for ( mlvl_bboxes , mlvl_scores ,
mlvl_conf_scores ) in zip ( batch_mlvl_bboxes , batch_mlvl_scores ,
batch_mlvl_conf_scores ) :
# Filtering out all predictions with conf < conf_thr
conf_thr = cfg . get ( ' conf_thr ' , - 1 )
if conf_thr > 0 and ( not torch . onnx . is_in_onnx_export ( ) ) :
# TensorRT not support NonZero
# add as_tuple=False for compatibility in Pytorch 1.6
# flatten would create a Reshape op with constant values,
# and raise RuntimeError when doing inference in ONNX
# Runtime with a different input image (#4221).
conf_inds = mlvl_conf_scores . ge ( conf_thr ) . nonzero (
as_tuple = False ) . squeeze ( 1 )
mlvl_bboxes = mlvl_bboxes [ conf_inds , : ]
mlvl_scores = mlvl_scores [ conf_inds , : ]
mlvl_conf_scores = mlvl_conf_scores [ conf_inds ]
det_bboxes , det_labels = multiclass_nms (
mlvl_bboxes ,
mlvl_scores ,
cfg . score_thr ,
cfg . nms ,
cfg . max_per_img ,
score_factors = mlvl_conf_scores )
det_results . append ( tuple ( [ det_bboxes , det_labels ] ) )
else :
return ( multi_lvl_bboxes , multi_lvl_cls_scores ,
multi_lvl_conf_scores )
det_results = [
tuple ( mlvl_bs )
for mlvl_bs in zip ( batch_mlvl_bboxes , batch_mlvl_scores ,
batch_mlvl_conf_scores )
]
return det_results
@force_fp32 ( apply_to = ( ' pred_maps ' , ) )
def loss ( self ,
Haian Huang(深度眸)
4 years ago
committed by
GitHub