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