def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
#print("blob_conv,",[x.shape for x in blob_conv])
rpn_ret = self.RPN(blob_conv, im_info, roidb)
# if self.training:
# # can be used to infer fg/bg ratio
# return_dict['rois_label'] = rpn_ret['labels_int32']
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
#print(box_feat.shape, res5_feat.shape)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
#print("main_concept", box_feat.shape, bbox_pred.shape, )
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(
dict((k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'],
rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'],
rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(
res5_feat,
rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(
mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(
res5_feat,
rpn_ret,
roi_has_keypoints_int32=rpn_ret[
'roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
# if self.training:
# # can be used to infer fg/bg ratio
# return_dict['rois_label'] = rpn_ret['labels_int32']
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(dict(
(k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits') or k.startswith('rpn_bbox_pred'))
))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(range(cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' % lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' % lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'], rpn_ret['bbox_targets'],
rpn_ret['bbox_inside_weights'], rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(res5_feat, rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(res5_feat, rpn_ret,
roi_has_keypoints_int32=rpn_ret['roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'], rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'], rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
# roidb: list, length = batch size
# 'has_visible_keypoints': bool
# 'boxes' & 'gt_classes': object bboxes and classes
# 'segms', 'seg_areas', 'gt_overlaps', 'is_crowd', 'box_to_gt_ind_map': pass
# 'gt_actions': num_box*26
# 'gt_role_id': num_box*26*2
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
# Original RPN module will generate proposals, and sample 256 positive/negative
# examples in a 1:3 ratio for r-cnn stage. For InteractNet(hoi), I set
# cfg.TRAIN.BATCH_SIZE_PER_IM and cfg.TRAIN.FG_FRACTION big value, to save
# every proposal in rpn_ret, then I will re-sample from rpn_ret for three branch
# of InteractNet, see roi_data/hoi_data.py for more information.
if not cfg.VCOCO.USE_PRECOMP_BOX:
rpn_ret = self.RPN(blob_conv, im_info, roidb)
if cfg.MODEL.VCOCO_ON and self.training:
# WARNING! always sample hoi branch before detection branch when training
hoi_blob_in = sample_for_hoi_branch(rpn_ret,
roidb,
im_info,
is_training=True)
# Re-sampling for RCNN head, rpn_ret will be modified inplace
sample_for_detection_branch(rpn_ret)
elif self.training:
json_dataset.add_proposals(roidb,
rois=None,
im_info=im_info.data.numpy(),
crowd_thresh=0) #[:, 2]
hoi_blob_in = sample_for_hoi_branch_precomp_box_train(
roidb, im_info, is_training=True)
if hoi_blob_in is None:
return_dict['losses'] = {}
return_dict['metrics'] = {}
return_dict['losses'][
'loss_hoi_interaction_action'] = torch.tensor(
[0.]).cuda(device_id)
return_dict['metrics'][
'accuracy_interaction_cls'] = torch.tensor(
[0.]).cuda(device_id)
return_dict['losses'][
'loss_hoi_interaction_affinity'] = torch.tensor(
[0.]).cuda(device_id)
return_dict['metrics'][
'accuracy_interaction_affinity'] = torch.tensor(
[0.]).cuda(device_id)
return return_dict
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
#blob_conv = blob_conv[-self.num_roi_levels:]
if cfg.FPN.MULTILEVEL_ROIS:
blob_conv = blob_conv[-self.num_roi_levels:]
else:
blob_conv = blob_conv[-1]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.VCOCO.USE_PRECOMP_BOX:
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
if not cfg.VCOCO.USE_PRECOMP_BOX:
rpn_kwargs.update(
dict((k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'],
rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'],
rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(
res5_feat,
rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(
mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(
res5_feat,
rpn_ret,
roi_has_keypoints_int32=rpn_ret[
'roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
if cfg.MODEL.VCOCO_ON:
hoi_blob_out = self.HOI_Head(blob_conv, hoi_blob_in)
interaction_action_loss, interaction_affinity_loss, \
interaction_action_accuracy_cls, interaction_affinity_cls = self.HOI_Head.loss(
hoi_blob_out)
return_dict['losses'][
'loss_hoi_interaction_action'] = interaction_action_loss
return_dict['metrics'][
'accuracy_interaction_cls'] = interaction_action_accuracy_cls
return_dict['losses'][
'loss_hoi_interaction_affinity'] = interaction_affinity_loss
return_dict['metrics'][
'accuracy_interaction_affinity'] = interaction_affinity_cls
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
else:
if not cfg.VCOCO.USE_PRECOMP_BOX:
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
#print('return ready')
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
# if self.training:
# # can be used to infer fg/bg ratio
# return_dict['rois_label'] = rpn_ret['labels_int32']
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(
dict((k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for k, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] = loss_rpn_cls[k]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] = loss_rpn_bbox[k]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score,
bbox_pred,
rpn_ret['labels_int32'],
rpn_ret['bbox_targets'],
rpn_ret['bbox_inside_weights'],
rpn_ret['bbox_outside_weights'],
stage=0)
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(
res5_feat,
rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(
mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(
res5_feat,
rpn_ret,
roi_has_keypoints_int32=rpn_ret[
'roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
if not cfg.FAST_RCNN.USE_CASCADE:
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
else:
return_dict['rois' + '_{}'.format(0)] = rpn_ret['rois']
return_dict['cls_score' + '_{}'.format(0)] = cls_score.detach()
return_dict['bbox_pred' + '_{}'.format(0)] = bbox_pred.detach()
if cfg.FAST_RCNN.USE_CASCADE:
for i in range(2):
i += 1
pre_stage_name = '_{}'.format(i - 1)
cls_score_cpu = cls_score.data.cpu().numpy()
import utils.boxes as box_utils
bbox_pred_cpu = bbox_pred.reshape(
[-1, bbox_pred.shape[-1]]).data.cpu().numpy().squeeze()
rois = deepcopy(return_dict['rois' + pre_stage_name])
assert cfg.MODEL.CLS_AGNOSTIC_BBOX_REG is True
if not cfg.MODEL.CLS_AGNOSTIC_BBOX_REG:
cls_loc = np.argmax(cls_score_cpu[:, 1:], axis=1) + 1
cls_loc = np.reshape(cls_loc, (cls_loc.shape[0], 1))
# Based on scores, we can select transformed rois
scores = np.zeros((cls_score_cpu.shape[0], 1))
for k in range(len(cls_loc)):
scores[k] = cls_score_cpu[k, cls_loc[k]]
batch_inds = rois[:, 0]
uni_inds = np.unique(batch_inds)
# WE suppose the WIDTH of image is EQUAL to its HEIGHT
# We also provide an example to show how to perform the operation batch-wise
# Scale forward
batch_se = []
for e in range(len(uni_inds)):
id_min = min(np.where(batch_inds == uni_inds[e])[0])
id_max = max(np.where(batch_inds == uni_inds[e])[0])
rois[id_min:id_max + 1, 1:5] /= im_info[e][2]
batch_se.append([id_min, id_max])
pred_boxes = box_utils.bbox_transform(
rois[:, 1:5], bbox_pred_cpu,
cfg.CASCADE_RCNN.BBOX_REG_WEIGHTS[i])
# Scale back
for e in range(len(uni_inds)):
id_min = batch_se[e][0]
id_max = batch_se[e][1]
pred_boxes[id_min:id_max + 1] *= im_info[e][2]
pred_boxes[id_min:id_max + 1] = box_utils.clip_tiled_boxes(
pred_boxes[id_min:id_max + 1], im_info[e][0:2])
cfg_key = 'TRAIN' if self.training else 'TEST'
min_size = cfg[cfg_key].RPN_MIN_SIZE
if not cfg.MODEL.CLS_AGNOSTIC_BBOX_REG:
# Cannot use for loop here which may cause "illegal memory access"
# Thanks to Chen-Wei Xie !
rows = pred_boxes.shape[0]
b3 = cls_loc * 4 + np.array([0, 1, 2, 3])
b4 = np.array(range(rows))
c = pred_boxes[np.repeat(b4, 4), b3.flatten()]
proposals = np.reshape(c, (-1, 4))
else:
proposals = pred_boxes[:, 4:8]
keep = _filter_boxes(proposals, min_size, im_info[0])
proposals = proposals[keep, :]
batch_inds = batch_inds[keep]
batch_inds = np.reshape(batch_inds, [len(batch_inds), 1])
proposals = np.concatenate((batch_inds, proposals), axis=1)
from modeling.collect_and_distribute_fpn_rpn_proposals import CollectAndDistributeFpnRpnProposalsOp
self.CollectAndDistributeFpnRpnProposals = CollectAndDistributeFpnRpnProposalsOp(
)
self.CollectAndDistributeFpnRpnProposals.training = self.training
# proposals.astype('float32')
blobs_out = self.CollectAndDistributeFpnRpnProposals(proposals,
roidb,
im_info,
stage=i)
# Update rpn_ret
new_rpn_ret = {}
for key, value in rpn_ret.items():
if 'rpn' in key:
new_rpn_ret[key] = value
new_rpn_ret.update(blobs_out)
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if i == 1:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head_2(
blob_conv, new_rpn_ret)
else:
box_feat = self.Box_Head_2(blob_conv, new_rpn_ret)
cls_score, bbox_pred = self.Box_Outs_2(box_feat)
elif i == 2:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head_3(
blob_conv, new_rpn_ret)
else:
box_feat = self.Box_Head_3(blob_conv, new_rpn_ret)
cls_score, bbox_pred = self.Box_Outs_3(box_feat)
if self.training:
# rpn loss
rpn_kwargs.update(
dict((k, new_rpn_ret[k]) for k in new_rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for k, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] += loss_rpn_cls[k]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] += loss_rpn_bbox[k]
else:
return_dict['losses']['loss_rpn_cls'] += loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] += loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score,
bbox_pred,
new_rpn_ret['labels_int32'],
new_rpn_ret['bbox_targets'],
new_rpn_ret['bbox_inside_weights'],
new_rpn_ret['bbox_outside_weights'],
stage=i)
return_dict['losses']['loss_cls'] += loss_cls
return_dict['losses']['loss_bbox'] += loss_bbox
return_dict['metrics']['accuracy_cls'] += accuracy_cls
return_dict['rois' + '_{}'.format(i)] = deepcopy(
new_rpn_ret['rois'])
return_dict['cls_score' + '_{}'.format(i)] = cls_score.detach()
return_dict['bbox_pred' + '_{}'.format(i)] = bbox_pred.detach()
rpn_ret = new_rpn_ret.copy()
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
# rpn proposals
# if self.training:
# # can be used to infer fg/bg ratio
# return_dict['rois_label'] = rpn_ret['labels_int32']
rois_certification = False
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if rois_certification:
lvl_min = cfg.FPN.ROI_MIN_LEVEL
lvl_max = cfg.FPN.ROI_MAX_LEVEL
test_rpn_ret = {'rois': rpn_ret['rois']}
lvls = fpn_utils.map_rois_to_fpn_levels(test_rpn_ret['rois'],
lvl_min, lvl_max)
rois_idx_order = np.empty((0, ))
test_rois = test_rpn_ret['rois']
for output_idx, lvl in enumerate(range(lvl_min, lvl_max + 1)):
idx_lvl = np.where(lvls == lvl)[0]
rois_lvl = test_rois[idx_lvl, :]
rois_idx_order = np.concatenate((rois_idx_order, idx_lvl))
test_rpn_ret['rois_fpn{}'.format(lvl)] = rois_lvl
rois_idx_restore = np.argsort(rois_idx_order).astype(np.int32,
copy=False)
test_rpn_ret['rois_idx_restore_int32'] = rois_idx_restore
test_feat = self.Box_Head(blob_conv, test_rpn_ret)
test_cls_score, test_bbox_pred = self.Box_Outs(test_feat)
test_cls_score = test_cls_score.data.cpu().numpy().squeeze()
test_bbox_pred = test_bbox_pred.data.cpu().numpy().squeeze()
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
# bbox proposals
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
# 在这里开始计算loss
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(
dict((k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'],
rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'],
rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(
res5_feat,
rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(
mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(
res5_feat,
rpn_ret,
roi_has_keypoints_int32=rpn_ret[
'roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
import json
if cfg.TEST.IOU_OUT:
# 直接通过rpn_ret可以取出rois
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/raw_roi.json",
'w') as f:
json.dump((return_dict['rois'][:, 1:] /
im_info.numpy()[0][2]).tolist(), f)
# 如果在FPN模式下，需要进到一个collect_and_distribute...的函数去取出分配后的scores
# ，我直接在collect_and_distribute_fpn_rpn_proposals.py里把json输出
# 因此这里直接考虑RPN_ONLY模式的取值。
if not cfg.FPN.FPN_ON:
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/rois_score.json",
'w') as f:
score_2_json = []
for item in rpn_ret['rpn_roi_probs']:
score_2_json.append(item.item())
json.dump(score_2_json, f)
# 开始第二个fast_head阶段，首先通过rois和bbox_delta计算pred_box
if cfg.FAST_RCNN.FAST_HEAD2_DEBUG:
lvl_min = cfg.FPN.ROI_MIN_LEVEL
lvl_max = cfg.FPN.ROI_MAX_LEVEL
if cfg.FPN.FPN_ON:
im_scale = im_info.data.cpu().numpy().squeeze()[2]
rois = rpn_ret['rois'][:, 1:5] / im_scale
bbox_pred = bbox_pred.data.cpu().numpy().squeeze()
box_deltas = bbox_pred.reshape([-1, bbox_pred.shape[-1]])
shift_boxes = box_utils.bbox_transform(
rois, box_deltas, cfg.MODEL.BBOX_REG_WEIGHTS)
shift_boxes = box_utils.clip_tiled_boxes(
shift_boxes,
im_info.data.cpu().numpy().squeeze()[0:2])
num_classes = cfg.MODEL.NUM_CLASSES
onecls_pred_boxes = []
onecls_score = []
dets_cls = {}
count = 0
for j in range(1, num_classes):
inds = np.where(
cls_score[:, j] > cfg.TEST.SCORE_THRESH)[0]
boxes_j = shift_boxes[inds, j * 4:(j + 1) * 4]
score_j = cls_score[inds, j]
onecls_pred_boxes += boxes_j.tolist()
onecls_score += score_j.tolist()
dets_cls.update({j: []})
for k in range(len(boxes_j.tolist())):
dets_cls[j].append(count)
count += 1
assert count == len(onecls_pred_boxes)
stage2_rois_score = np.array(onecls_score,
dtype=np.float32)
stage2_rois = np.array(onecls_pred_boxes, dtype=np.float32)
# Redistribute stage2_rois using fpn_utils module provided functions
# calculate by formula
cls_tracker = {}
if not stage2_rois.tolist():
stage1_pred_iou = stage2_rois_score.tolist()
stage2_final_boxes = np.empty((0, ))
stage2_final_score = np.empty((0, ))
logger.info("Detections above threshold is null.")
else:
alter_rpn = {}
unresize_stage2_rois = stage2_rois * im_scale
# unresize_stage2_rois = np.concatenate((unresize_stage2_rois, unresized_rois[:, 1:5]))
lvls = fpn_utils.map_rois_to_fpn_levels(
unresize_stage2_rois, lvl_min, lvl_max)
# TAG: We might need to visualize "stage2_rois" to make sure.
rois_idx_order = np.empty((0, ))
dummy_batch = np.zeros(
(unresize_stage2_rois.shape[0], 1),
dtype=np.float32)
alter_rpn["rois"] = np.hstack(
(dummy_batch,
unresize_stage2_rois)).astype(np.float32,
copy=False)
# alter_rpn['rois'] = np.concatenate((alter_rpn['rois'], unresized_rois))
for output_idx, lvl in enumerate(
range(lvl_min, lvl_max + 1)):
idx_lvl = np.where(lvls == lvl)[0]
rois_lvl = unresize_stage2_rois[idx_lvl, :]
rois_idx_order = np.concatenate(
(rois_idx_order, idx_lvl))
_ = np.zeros((rois_lvl.shape[0], 1),
dtype=np.float32)
alter_rpn['rois_fpn{}'.format(lvl)] = np.hstack(
(_, rois_lvl)).astype(np.float32, copy=False)
rois_idx_restore = np.argsort(rois_idx_order).astype(
np.int32, copy=False)
alter_rpn['rois_idx_restore_int32'] = rois_idx_restore
# Go through 2nd stage of FPN and fast_head
stage2_feat = self.Box_Head(blob_conv, alter_rpn)
stage2_cls_score, stage2_bbox_pred = self.Box_Outs(
stage2_feat)
# Transform shift value to original one to get final pred boxes coordinates
stage2_bbox_pred = stage2_bbox_pred.data.cpu().numpy(
).squeeze()
stage2_cls_score = stage2_cls_score.data.cpu().numpy()
stage2_box_deltas = stage2_bbox_pred.reshape(
[-1, bbox_pred.shape[-1]])
# Add some variance to box delta
if cfg.FAST_RCNN.STAGE1_TURBULENCE:
import random
for i in range(len(stage2_box_deltas)):
for j in range(len(stage2_box_deltas[i])):
stage2_box_deltas[i][j] *= random.uniform(
0.9, 1.1)
stage2_cls_out = box_utils.bbox_transform(
stage2_rois, stage2_box_deltas,
cfg.MODEL.BBOX_REG_WEIGHTS)
stage2_cls_out = box_utils.clip_tiled_boxes(
stage2_cls_out,
im_info.data.cpu().numpy().squeeze()[0:2])
onecls_pred_boxes = []
onecls_score = []
count = 0
for j in range(1, num_classes):
inds = np.where(
stage2_cls_score[:,
j] > cfg.TEST.SCORE_THRESH)[0]
boxes_j = stage2_cls_out[inds, j * 4:(j + 1) * 4]
score_j = stage2_cls_score[inds, j]
dets_j = np.hstack(
(boxes_j,
score_j[:, np.newaxis])).astype(np.float32,
copy=False)
keep = box_utils.nms(dets_j, cfg.TEST.NMS)
boxes_j = boxes_j[keep]
score_j = score_j[keep]
# 用于记录每个框属于第几类
onecls_score += score_j.tolist()
onecls_pred_boxes += boxes_j.tolist()
for k in range(len(score_j)):
cls_tracker.update({count: j})
count += 1
assert count == len(onecls_score)
stage2_final_boxes = np.array(onecls_pred_boxes,
dtype=np.float32)
stage2_final_score = np.array(onecls_score,
dtype=np.float32)
inds = np.where(stage2_final_score > 0.3)[0]
# Filtered by keep index...
preserve_stage2_final_boxes = copy.deepcopy(
stage2_final_boxes)
preserve_stage2_final_score = copy.deepcopy(
stage2_final_score)
stage2_final_boxes = stage2_final_boxes[inds]
stage2_final_score = stage2_final_score[inds]
# if nothing left after 0.3 threshold filter, reserve whole boxes to original.
if stage2_final_boxes.size == 0:
lower_inds = np.where(
preserve_stage2_final_score > 0.1)[0]
stage2_final_boxes = preserve_stage2_final_boxes[
lower_inds]
stage2_final_score = preserve_stage2_final_score[
lower_inds]
else:
del preserve_stage2_final_boxes
del preserve_stage2_final_score
# if all boxes are clsfied into bg class.
if stage2_final_boxes.size == 0:
stage1_pred_iou = stage2_rois_score.tolist()
stage2_final_boxes = np.empty((0, ))
stage2_final_score = np.empty((0, ))
logger.info("Detections above threshold is null.")
else:
# Restore stage2_pred_boxes to match the index with stage2_rois, Compute IOU between
# final_boxes and stage2_rois, one by one
flag = "cross_product"
if flag == "element_wise":
if stage2_final_boxes.shape[
0] == stage2_rois.shape[0]:
restored_stage2_final_boxes = stage2_final_boxes[
rois_idx_restore]
stage1_pred_iou = []
for ind, item in enumerate(stage2_rois):
stage1 = np.array(
item, dtype=np.float32).reshape(
(1, 4))
stage2 = np.array(
restored_stage2_final_boxes[ind],
dtype=np.float32).reshape((1, 4))
iou = box_utils.bbox_overlaps(
stage1, stage2)
stage1_pred_iou.append(
iou.squeeze().item())
else:
logger.info(
"Mistake while processing {}".format(
str(im_info)))
elif flag == "cross_product":
iou = box_utils.bbox_overlaps(
stage2_rois, stage2_final_boxes)
stage1_pred_iou = iou.max(axis=1).tolist()
# stage1_pred is another name of stage2_rois
assert len(stage1_pred_iou) == len(stage2_rois)
if cfg.FAST_RCNN.IOU_NMS:
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_stage1_score.json",
"w") as f:
json.dump(stage2_rois_score.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_stage2_score.json",
"w") as f:
json.dump(stage2_final_score.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_stage1_pred_boxes.json",
'w') as f:
json.dump(stage2_rois.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_stage1_pred_iou.json",
'w') as f:
json.dump(stage1_pred_iou, f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_stage2_pred_boxes.json",
'w') as f:
json.dump(stage2_final_boxes.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_dets_cls.json",
'w') as f:
json.dump(dets_cls, f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/iou_cls_tracker.json",
'w') as f:
json.dump(cls_tracker, f)
elif cfg.FAST_RCNN.SCORE_NMS:
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_stage1_score.json",
"w") as f:
json.dump(stage2_rois_score.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_stage2_score.json",
"w") as f:
json.dump(stage2_final_score.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_stage1_pred_boxes.json",
'w') as f:
json.dump(stage2_rois.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_stage1_pred_iou.json",
'w') as f:
json.dump(stage1_pred_iou, f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_stage2_pred_boxes.json",
'w') as f:
json.dump(stage2_final_boxes.tolist(), f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_dets_cls.json",
'w') as f:
json.dump(dets_cls, f)
with open(
"/nfs/project/libo_i/IOU.pytorch/IOU_Validation/score_cls_tracker.json",
'w') as f:
json.dump(cls_tracker, f)
else:
im_scale = im_info.data.cpu().numpy().squeeze()[2]
rois = rpn_ret['rois'][:, 1:5] / im_scale
# unscale back to raw image space
box_deltas = bbox_pred.data.cpu().numpy().squeeze()
fast_stage1_score = cls_score.data.cpu().numpy().squeeze()
box_deltas = box_deltas.reshape([-1, bbox_pred.shape[-1]])
stage2_rois = box_utils.bbox_transform(
rois, box_deltas, cfg.MODEL.BBOX_REG_WEIGHTS)
stage2_rois = box_utils.clip_tiled_boxes(
stage2_rois,
im_info.data.cpu().numpy().squeeze()[0:2])
num_classes = cfg.MODEL.NUM_CLASSES
onecls_pred_boxes = []
onecls_cls_score = []
for j in range(1, num_classes):
inds = np.where(
cls_score[:, j] > cfg.TEST.SCORE_THRESH)[0]
boxes_j = stage2_rois[inds, j * 4:(j + 1) * 4]
score_j = fast_stage1_score[inds, j]
onecls_pred_boxes += boxes_j.tolist()
onecls_cls_score += score_j.tolist()
stage2_rois = np.array(onecls_pred_boxes, dtype=np.float32)
stage2_rois_score = np.array(onecls_cls_score,
dtype=np.float32)
assert len(stage2_rois) == len(stage2_rois_score)
# Send stage2 rois to next stage fast head, do ROI ALIGN again
# to modify rpn_ret['rois] , rpn_ret['rpn_rois'] and rpn['rois_rpn_score']
rpn_ret['rois'] = stage2_rois
rpn_ret['rpn_rois'] = stage2_rois
rpn_ret['rpn_roi_probs'] = stage2_rois_score
stage2_box_feat = self.Box_Head(blob_conv, rpn_ret)
stage2_cls_score, stage2_bbox_pred = self.Box_Outs(
stage2_box_feat)
stage2_bbox_pred = stage2_bbox_pred.data.cpu().numpy(
).squeeze()
stage2_bbox_pred = stage2_bbox_pred.reshape(
[-1, bbox_pred.shape[-1]])
stage2_cls_pred_boxes = box_utils.bbox_transform(
stage2_rois, stage2_bbox_pred,
cfg.MODEL.BBOX_REG_WEIGHTS)
stage2_cls_pred_boxes = box_utils.clip_tiled_boxes(
stage2_cls_pred_boxes,
im_info.data.cpu().numpy().squeeze()[0:2])
onecls_pred_boxes = []
onecls_cls_score = []
for j in range(1, num_classes):
inds = np.where(
stage2_cls_score[:, j] > cfg.TEST.SCORE_THRESH)[0]
if len(inds) != 0:
print("KKKKK")
boxes_j = stage2_cls_pred_boxes[inds,
j * 4:(j + 1) * 4]
score_j = stage2_cls_score[inds, j]
onecls_pred_boxes += boxes_j.tolist()
onecls_cls_score += score_j.tolist()
stage2_bbox_pred = np.array(onecls_pred_boxes,
dtype=np.float32)
stage2_bbox_pred_score = np.array(onecls_cls_score,
dtype=np.float32)
# get stage2 pred_boxes here
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
# if self.training:
# # can be used to infer fg/bg ratio
# return_dict['rois_label'] = rpn_ret['labels_int32']
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(
dict((k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'],
rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'],
rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.TRANS_HEAD_ON:
# We first get the bounding box
pred_boxes = car_3d_pose_heads.bbox_transform_pytorch(
rpn_ret['rois'], bbox_pred, cfg.MODEL.BBOX_REG_WEIGHTS)
# we only use the car cls
car_cls_int = 4
car_idx = np.where(rpn_ret['labels_int32'] == car_cls_int)
pred_boxes_car = pred_boxes[car_idx, 4 * car_cls_int:4 *
(car_cls_int + 1)].squeeze(dim=0)
# Build translation head heres from the bounding box
car_trans_feat = self.car_trans_Head(pred_boxes_car)
car_trans_pred = self.car_trans_Outs(car_trans_feat)
label_trans = rpn_ret['car_trans'][car_idx]
loss_trans = car_3d_pose_heads.car_trans_losses(
car_trans_pred, label_trans)
return_dict['losses']['loss_trans'] = loss_trans
if cfg.MODEL.CAR_CLS_HEAD_ON:
if getattr(self.car_cls_Head, 'SHARE_RES5', False):
# TODO: add thos shared_res5 module
pass
else:
car_cls_rot_feat = self.car_cls_Head(blob_conv, rpn_ret)
car_cls_score, car_cls, rot_pred = self.car_cls_Outs(
car_cls_rot_feat)
# car classification loss, we only fine tune the labelled cars
# we only use the car cls
car_idx = np.where(rpn_ret['labels_int32'] == 4)
if len(cfg.TRAIN.CE_CAR_CLS_FINETUNE_WIGHT):
ce_weight = np.array(cfg.TRAIN.CE_CAR_CLS_FINETUNE_WIGHT)
else:
ce_weight = None
shape_sim_mat = self.shape_sim_mat_loss[
rpn_ret['car_cls_labels_int32'][car_idx].astype('int64')]
if cfg.CAR_CLS.SIM_MAT_LOSS:
shape_sim_mat_loss_mat = []
#print('Not implemented properly yet')
else:
shape_sim_mat_loss_mat = []
loss_car_cls, loss_rot, accuracy_car_cls = car_3d_pose_heads.fast_rcnn_car_cls_rot_losses(
car_cls_score[car_idx],
rot_pred[car_idx],
car_cls[car_idx],
rpn_ret['car_cls_labels_int32'][car_idx],
rpn_ret['quaternions'][car_idx],
ce_weight,
shape_sim_mat_loss_mat=shape_sim_mat_loss_mat)
return_dict['losses']['loss_car_cls'] = loss_car_cls
return_dict['losses']['loss_rot'] = loss_rot
return_dict['metrics']['accuracy_car_cls'] = accuracy_car_cls
return_dict['metrics']['shape_sim'] = shape_sim(
car_cls[car_idx].data.cpu().numpy(), shape_sim_mat)
return_dict['metrics']['rot_sim'] = rot_sim(
rot_pred[car_idx].data.cpu().numpy(),
rpn_ret['quaternions'][car_idx])
return_dict['metrics']['trans_sim'] = trans_sim(
car_trans_pred.data.cpu().numpy(),
rpn_ret['car_trans'][car_idx], cfg.TRANS_HEAD.TRANS_MEAN,
cfg.TRANS_HEAD.TRANS_STD)
if cfg.MODEL.MASK_TRAIN_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(
res5_feat,
rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(
mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(
res5_feat,
rpn_ret,
roi_has_keypoints_int32=rpn_ret[
'roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
if type(v) == np.float64:
return_dict['metrics'][k] = v
else:
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
return_dict = {} # A dict to collect return variables
if cfg.FPN.NON_LOCAL:
blob_conv, f_div_C = self.Conv_Body(im_data)
if cfg.MODEL.NON_LOCAL_TEST:
return_dict['f_div_C'] = f_div_C
else:
blob_conv = self.Conv_Body(im_data)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(
dict((k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits')
or k.startswith('rpn_bbox_pred'))))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(
**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(
range(cfg.FPN.RPN_MIN_LEVEL,
cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' %
lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' %
lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'],
rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'],
rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
# we only use the car cls
if np.sum(rpn_ret['labels_int32']) <= 0:
print("ZERO POSITIVE")
if cfg.MODEL.CAR_CLS_HEAD_ON:
if getattr(self.car_cls_Head, 'SHARE_RES5', False):
# TODO: add thos shared_res5 module
pass
else:
car_cls_rot_feat = self.car_cls_Head(blob_conv, rpn_ret)
car_cls_score, car_cls, rot_pred = self.car_cls_Outs(
car_cls_rot_feat)
# car classification loss, we only fine tune the labelled cars
# we only use the car cls
care_idx = np.where(
list(
map(lambda x: x in cfg.TRAIN.CARE_CLS,
rpn_ret['labels_int32'])))
if len(cfg.TRAIN.CE_CAR_CLS_FINETUNE_WIGHT):
ce_weight = np.array(cfg.TRAIN.CE_CAR_CLS_FINETUNE_WIGHT)
else:
ce_weight = []
loss_car_cls, loss_rot, accuracy_car_cls = car_3d_pose_heads.fast_rcnn_car_cls_rot_losses(
car_cls_score[care_idx],
rot_pred[care_idx],
car_cls[care_idx],
rpn_ret['car_cls_labels_int32'][care_idx],
rpn_ret['quaternions'][care_idx],
ce_weight,
shape_sim_mat=self.shape_sim_mat)
return_dict['losses']['loss_car_cls'] = loss_car_cls
return_dict['losses']['loss_rot'] = loss_rot
if cfg.CAR_CLS.CLS_LOSS:
return_dict['metrics'][
'accuracy_car_cls'] = accuracy_car_cls
return_dict['metrics']['shape_sim'] = shape_sim(
car_cls[care_idx].data.cpu().numpy(),
self.shape_sim_mat, rpn_ret['car_cls_labels_int32']
[care_idx].astype('int64'))
return_dict['metrics']['rot_diff_degree'] = rot_sim(
rot_pred[care_idx].data.cpu().numpy(),
rpn_ret['quaternions'][care_idx])
if cfg.MODEL.TRANS_HEAD_ON:
pred_boxes = car_3d_pose_heads.bbox_transform_pytorch(
rpn_ret['rois'], bbox_pred, im_info,
cfg.MODEL.BBOX_REG_WEIGHTS)
care_idx = np.where(
list(
map(lambda x: x in cfg.TRAIN.CARE_CLS,
rpn_ret['labels_int32'])))
# Build translation head heres from the bounding box
if cfg.TRANS_HEAD.INPUT_CONV_BODY:
pred_boxes_tmp = [
pred_boxes[i, 4 * rpn_ret['labels_int32'][i]:4 *
(rpn_ret['labels_int32'][i] + 1)]
for i in range(rpn_ret['labels_int32'].shape[0])
]
pred_boxes_car = torch.stack(pred_boxes_tmp).squeeze(dim=0)
car_trans_feat = self.car_trans_Head(
blob_conv, rpn_ret, pred_boxes_car)
car_trans_pred = self.car_trans_Outs(car_trans_feat)
car_trans_pred = car_trans_pred[care_idx]
elif cfg.TRANS_HEAD.INPUT_TRIPLE_HEAD:
pred_boxes_tmp = [
pred_boxes[i, 4 * rpn_ret['labels_int32'][i]:4 *
(rpn_ret['labels_int32'][i] + 1)]
for i in range(rpn_ret['labels_int32'].shape[0])
]
pred_boxes_car = torch.stack(pred_boxes_tmp).squeeze(dim=0)
car_trans_feat = self.car_trans_Head(pred_boxes_car)
car_trans_pred = self.car_trans_Outs(
car_trans_feat, car_cls_rot_feat)
car_trans_pred = car_trans_pred[care_idx]
else:
car_cls_int = 4
pred_boxes_car = pred_boxes[care_idx, 4 * car_cls_int:4 *
(car_cls_int + 1)].squeeze(
dim=0)
car_trans_feat = self.car_trans_Head(pred_boxes_car)
car_trans_pred = self.car_trans_Outs(car_trans_feat)
label_trans = rpn_ret['car_trans'][care_idx]
if cfg.MODEL.Z_MEAN > 0:
label_trans = label_trans / 100
#label_trans[:, -1] = label_trans[:, -1] - 600
loss_trans = car_3d_pose_heads.car_trans_losses(
car_trans_pred, label_trans)
return_dict['losses']['loss_trans'] = loss_trans
return_dict['metrics']['trans_diff_meter'], return_dict['metrics']['trans_thresh_per'] = \
trans_sim(car_trans_pred.data.cpu().numpy(), rpn_ret['car_trans'][care_idx],
cfg.TRANS_HEAD.TRANS_MEAN, cfg.TRANS_HEAD.TRANS_STD)
# A 3D to 2D projection loss
if cfg.MODEL.LOSS_3D_2D_ON:
# During the mesh generation, using GT(True) or predicted(False) Car ID
if cfg.LOSS_3D_2D.MESH_GEN_USING_GT:
# Acquire car id
car_ids = rpn_ret['car_cls_labels_int32'][care_idx].astype(
'int64')
else:
# Using the predicted car id
print("Not properly implemented for pytorch")
car_ids = car_cls_score[care_idx].max(dim=1)
# Get mesh vertices and generate loss
UV_projection_loss = plane_projection_loss(
car_trans_pred, label_trans, rot_pred[care_idx],
rpn_ret['quaternions'][care_idx], car_ids, im_info,
self.car_models, self.intrinsic_mat, self.car_names)
return_dict['losses'][
'UV_projection_loss'] = UV_projection_loss
if cfg.MODEL.MASK_TRAIN_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(
res5_feat,
rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(
mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(
res5_feat,
rpn_ret,
roi_has_keypoints_int32=rpn_ret[
'roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'],
rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
if type(v) == np.float64:
return_dict['metrics'][k] = v
else:
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv = self.Conv_Body(im_data)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
# if self.training:
# # can be used to infer fg/bg ratio
# return_dict['rois_label'] = rpn_ret['labels_int32']
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
cls_score, bbox_pred = self.Box_Outs(box_feat)
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(dict(
(k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits') or k.startswith('rpn_bbox_pred'))
))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(range(cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' % lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' % lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'], rpn_ret['bbox_targets'],
rpn_ret['bbox_inside_weights'], rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
# EDIT: target soft-labels with distillation loss
if cfg.TRAIN.GT_SCORES:
# Smooth labels for the noisy dataset assumed to be "dataset-0"
if rpn_ret['dataset_id'][0] == 0:
if cfg.TRAIN.DISTILL_ATTN:
pre_lambda_attn = self.AttentionRoi_Head(box_feat)
loss_distill = fast_rcnn_heads.smooth_label_loss(
cls_score, rpn_ret['labels_int32'],
rpn_ret['gt_scores'],
rpn_ret['gt_source'],
dist_T=cfg.TRAIN.DISTILL_TEMPERATURE,
dist_lambda=F.sigmoid(pre_lambda_attn),
tracker_score=cfg.TRAIN.TRACKER_SCORE)
loss_distill += torch.sum(torch.abs(pre_lambda_attn)) # regularize attn magnitude
else:
loss_distill = fast_rcnn_heads.smooth_label_loss(
cls_score, rpn_ret['labels_int32'],
rpn_ret['gt_scores'],
rpn_ret['gt_source'],
dist_T=cfg.TRAIN.DISTILL_TEMPERATURE,
dist_lambda=cfg.TRAIN.DISTILL_LAMBDA,
tracker_score=cfg.TRAIN.TRACKER_SCORE)
return_dict['losses']['loss_cls'] = loss_distill
# EDIT: domain adversarial losses
if cfg.TRAIN.DOMAIN_ADAPT_IM or cfg.TRAIN.DOMAIN_ADAPT_ROI:
if rpn_ret['dataset_id'][0] == 0:
# dataset-0 is assumed unlabeled - so zero out other losses
# NOTE: multiplied by 0 to maintain valid autodiff graph
return_dict['losses']['loss_cls'] *= 0
return_dict['losses']['loss_bbox'] *= 0
return_dict['losses']['loss_rpn_cls'] *= 0
return_dict['losses']['loss_rpn_bbox'] *= 0
if cfg.TRAIN.DOMAIN_ADAPT_IM:
da_image_pred = self.DiscriminatorImage_Head(blob_conv)
loss_da_im = adversarial_heads.domain_loss_im(
da_image_pred, rpn_ret['dataset_id'][0])
return_dict['losses']['loss_da-im'] = loss_da_im
if cfg.TRAIN.DOMAIN_ADAPT_ROI:
da_roi_pred = self.DiscriminatorRoi_Head(box_feat)
loss_da_roi = adversarial_heads.domain_loss_im(
da_roi_pred, rpn_ret['dataset_id'][0])
return_dict['losses']['loss_da-roi'] = loss_da_roi
if cfg.TRAIN.DOMAIN_ADAPT_CST:
assert cfg.TRAIN.DOMAIN_ADAPT_ROI and cfg.TRAIN.DOMAIN_ADAPT_IM
loss_da_cst = adversarial_heads.domain_loss_cst(
da_image_pred, da_roi_pred)
return_dict['losses']['loss_da-cst'] = loss_da_cst
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(res5_feat, rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
if cfg.MODEL.KEYPOINTS_ON:
if getattr(self.Keypoint_Head, 'SHARE_RES5', False):
# No corresponding keypoint head implemented yet (Neither in Detectron)
# Also, rpn need to generate the label 'roi_has_keypoints_int32'
kps_feat = self.Keypoint_Head(res5_feat, rpn_ret,
roi_has_keypoints_int32=rpn_ret['roi_has_keypoint_int32'])
else:
kps_feat = self.Keypoint_Head(blob_conv, rpn_ret)
kps_pred = self.Keypoint_Outs(kps_feat)
# return_dict['keypoints_pred'] = kps_pred
# keypoints loss
if cfg.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'], rpn_ret['keypoint_weights'])
else:
loss_keypoints = keypoint_rcnn_heads.keypoint_losses(
kps_pred, rpn_ret['keypoint_locations_int32'], rpn_ret['keypoint_weights'],
rpn_ret['keypoint_loss_normalizer'])
return_dict['losses']['loss_kps'] = loss_keypoints
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict
def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
im_data = data
if self.training:
roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
device_id = im_data.get_device()
return_dict = {} # A dict to collect return variables
blob_conv,res5_feat = self.Conv_Body(im_data)
if cfg.MODEL.LR_VIEW_ON or cfg.MODEL.GIF_ON or cfg.MODEL.LRASY_MAHA_ON:
blob_conv = self._get_lrview_blob_conv(blob_conv)
rpn_ret = self.RPN(blob_conv, im_info, roidb)
if cfg.FPN.FPN_ON:
# Retain only the blobs that will be used for RoI heads. `blob_conv` may include
# extra blobs that are used for RPN proposals, but not for RoI heads.
blob_conv = blob_conv[-self.num_roi_levels:]
if not self.training:
return_dict['blob_conv'] = blob_conv
if not cfg.MODEL.RPN_ONLY:
if cfg.MODEL.SHARE_RES5 and self.training:
box_feat, res5_feat = self.Box_Head(blob_conv, rpn_ret)
context_box_feat, context_res5_feat = self.Context_Box_Head(blob_conv, rpn_ret)
else:
box_feat = self.Box_Head(blob_conv, rpn_ret)
context_box_feat = self.Context_Box_Head(blob_conv, rpn_ret)
if cfg.CONTEXT_ROI_POOLING.MERGE_TYPE == 'concat':
joint_box_feat = torch.cat([context_box_feat, box_feat], 1)
elif cfg.CONTEXT_ROI_POOLING.MERGE_TYPE == 'sum':
joint_box_feat = context_box_feat + box_feat
if cfg.CONTEXT_ROI_POOLING.SHARE_FEATURE:
cls_score, bbox_pred = self.Box_Outs(joint_box_feat, joint_box_feat)
else:
cls_score, bbox_pred = self.Box_Outs(joint_box_feat, box_feat)
# print cls_score.shape
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
if cfg.LESION.USE_POSITION:
position_feat = self.Position_Head(res5_feat)
pos_cls_pred = self.Position_Cls_Outs(position_feat)
pos_reg_pred = self.Position_Reg_Outs(position_feat)
return_dict['pos_cls_pred'] = pos_cls_pred
return_dict['pos_reg_pred'] = pos_reg_pred
else:
# TODO: complete the returns for RPN only situation
pass
if self.training:
return_dict['losses'] = {}
return_dict['metrics'] = {}
# rpn loss
rpn_kwargs.update(dict(
(k, rpn_ret[k]) for k in rpn_ret.keys()
if (k.startswith('rpn_cls_logits') or k.startswith('rpn_bbox_pred'))
))
loss_rpn_cls, loss_rpn_bbox = rpn_heads.generic_rpn_losses(**rpn_kwargs)
if cfg.FPN.FPN_ON:
for i, lvl in enumerate(range(cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL + 1)):
return_dict['losses']['loss_rpn_cls_fpn%d' % lvl] = loss_rpn_cls[i]
return_dict['losses']['loss_rpn_bbox_fpn%d' % lvl] = loss_rpn_bbox[i]
else:
return_dict['losses']['loss_rpn_cls'] = loss_rpn_cls
return_dict['losses']['loss_rpn_bbox'] = loss_rpn_bbox
if cfg.MODEL.FASTER_RCNN:
# bbox loss
loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
cls_score, bbox_pred, rpn_ret['labels_int32'], rpn_ret['bbox_targets'],
rpn_ret['bbox_inside_weights'], rpn_ret['bbox_outside_weights'])
return_dict['losses']['loss_cls'] = loss_cls
return_dict['losses']['loss_bbox'] = loss_bbox
return_dict['metrics']['accuracy_cls'] = accuracy_cls
if cfg.MODEL.MASK_ON:
if getattr(self.Mask_Head, 'SHARE_RES5', False):
mask_feat = self.Mask_Head(res5_feat, rpn_ret,
roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
else:
mask_feat = self.Mask_Head(blob_conv, rpn_ret)
mask_pred = self.Mask_Outs(mask_feat)
# return_dict['mask_pred'] = mask_pred
# mask loss
loss_mask = mask_rcnn_heads.mask_rcnn_losses(mask_pred, rpn_ret['masks_int32'])
return_dict['losses']['loss_mask'] = loss_mask
# pytorch0.4 bug on gathering scalar(0-dim) tensors
for k, v in return_dict['losses'].items():
return_dict['losses'][k] = v.unsqueeze(0)
for k, v in return_dict['metrics'].items():
return_dict['metrics'][k] = v.unsqueeze(0)
else:
# Testing
return_dict['rois'] = rpn_ret['rois']
return_dict['cls_score'] = cls_score
return_dict['bbox_pred'] = bbox_pred
return return_dict