def get_loss(cls_score, bbox_pred, blobs):
    if cls_score.size(0) != blobs['labels_int32'].shape[0]:
        pdb.set_trace()
    loss_cls, loss_bbox, accuracy_cls = fast_rcnn_heads.fast_rcnn_losses(
        cls_score, bbox_pred, blobs['labels_int32'], blobs['bbox_targets'],
        blobs['bbox_inside_weights'], blobs['bbox_outside_weights'])

    return loss_cls, loss_bbox, accuracy_cls
    def _forward(self, data, im_info, roidb=None, rois=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

        if cfg.MODEL.RELATION_NET_INPUT == 'GEO' and not cfg.REL_INFER.TRAIN and cfg.MODEL.NUM_RELATIONS > 0 and self.training:
            blob_conv = im_data
        else:
            blob_conv = self.Conv_Body(im_data)

        rpn_ret = self.RPN(blob_conv, im_info, roidb)
        # If rpn_ret doesn't have rpn_ret, the rois should have been given, we use that.
        if 'rois' not in rpn_ret:
            rpn_ret['rois'] = rois
        if hasattr(self, '_ignore_classes') and 'labels_int32' in rpn_ret:
            # Turn ignore classes labels to 0, because they are treated as background
            rpn_ret['labels_int32'][np.isin(rpn_ret['labels_int32'],
                                            self._ignore_classes)] = 0

        # 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:
                if cfg.MODEL.RELATION_NET_INPUT == 'GEO' and not cfg.REL_INFER.TRAIN and cfg.MODEL.NUM_RELATIONS > 0 and self.training:
                    box_feat = blob_conv.new_zeros(
                        rpn_ret['labels_int32'].shape[0],
                        self.Box_Head.dim_out)
                else:
                    box_feat = self.Box_Head(blob_conv, rpn_ret)
            if cfg.MODEL.RELATION_NET_INPUT == 'GEO' and not cfg.REL_INFER.TRAIN and cfg.MODEL.NUM_RELATIONS > 0 and self.training:
                cls_score = box_feat.new_zeros(
                    rpn_ret['labels_int32'].shape[0], cfg.MODEL.NUM_CLASSES)
                bbox_pred = box_feat.new_zeros(
                    rpn_ret['labels_int32'].shape[0],
                    4 * cfg.MODEL.NUM_CLASSES)
            else:
                cls_score, bbox_pred = self.Box_Outs(box_feat)
            if cfg.TEST.TAGGING:
                rois_label = torch.from_numpy(
                    rpn_ret['labels_int32'].astype('int64')).to(
                        cls_score.device)
                accuracy_cls = cls_score.max(1)[1].eq(rois_label).float().mean(
                    dim=0)
                print('Before refine:', accuracy_cls.item())

        if cfg.MODEL.NUM_RELATIONS > 0 and self.training:
            rel_scores, rel_labels = self.Rel_Outs(rpn_ret['rois'],
                                                   box_feat,
                                                   rpn_ret['labels_int32'],
                                                   roidb=roidb)
        elif cfg.MODEL.NUM_RELATIONS > 0 and cfg.TEST.USE_REL_INFER:
            if cfg.TEST.TAGGING:
                rel_scores = self.Rel_Outs(rpn_ret['rois'], box_feat)
                cls_score = self.Rel_Inf(cls_score, rel_scores, roidb=roidb)
            else:  # zero shot detection
                filt = (cls_score.topk(cfg.TEST.REL_INFER_PROPOSAL,
                                       1)[1] == 0).float().sum(1) == 0
                filt[cls_score[:,
                               1:].max(1)[0].topk(min(100, cls_score.shape[0]),
                                                  0)[1]] += 1
                filt = filt >= 2
                if filt.sum() == 0:
                    print('all background?')
                else:
                    tmp_rois = rpn_ret['rois'][filt.cpu().numpy().astype(
                        'bool')]

                    rel_scores = self.Rel_Outs(tmp_rois, box_feat[filt])
                    tmp_cls_score = self.Rel_Inf(cls_score[filt],
                                                 rel_scores,
                                                 roidb=None)
                    cls_score[filt] = tmp_cls_score
            if cfg.TEST.TAGGING:
                rois_label = torch.from_numpy(
                    rpn_ret['labels_int32'].astype('int64')).to(
                        cls_score.device)
                accuracy_cls = cls_score.max(1)[1].eq(rois_label).float().mean(
                    dim=0)
                print('After refine:', accuracy_cls.item())

        if self.training:
            return_dict['losses'] = {}
            return_dict['metrics'] = {}
            # rpn loss
            if not cfg.MODEL.TAGGING:
                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 not cfg.MODEL.RPN_ONLY:
                # 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 hasattr(loss_cls, 'mean_similarity'):
                    return_dict['metrics'][
                        'mean_similarity'] = loss_cls.mean_similarity

                if cfg.FAST_RCNN.SAE_REGU:
                    tmp = cls_score.sae_loss[(rpn_ret['labels_int32'] >
                                              0).tolist()]
                    return_dict['losses']['loss_sae'] = 1e-4 * tmp.mean(
                    ) if tmp.numel() > 0 else tmp.new_tensor(0)
                    if torch.isnan(return_dict['losses']['loss_sae']):
                        import pdb
                        pdb.set_trace()

            if cfg.REL_INFER.TRAIN:
                return_dict['losses']['loss_rel_infer'] = cfg.REL_INFER.TRAIN_WEIGHT * \
                    self.Rel_Inf_Train(rpn_ret['rois'], rpn_ret['labels_int32'], cls_score, rel_scores, roidb=roidb)

            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.update(rpn_ret)
            if not cfg.MODEL.RPN_ONLY:
                if cfg.TEST.KEEP_HIGHEST:
                    cls_score = F.softmax(cls_score * 1e10, dim=1) * cls_score
                return_dict['cls_score'] = cls_score
                return_dict['bbox_pred'] = bbox_pred

        return return_dict
示例#3
0
    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 = self.Conv_Body(im_data)  # (2, 1024, 50, 50), conv4

        rpn_ret = self.RPN(blob_conv, im_info, roidb)  # (N*2000, 5)

        # 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)  # box_feat: (N*512, 2014, 1, 1)
            else:
                # pdb.set_trace()
                box_feat = self.Box_Head(blob_conv, rpn_ret)
            cls_score, bbox_pred = self.Box_Outs(
                box_feat)  # cls_score: (N*512, C)
        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'))))

            # pdb.set_trace()
            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

            # LJY: if the image has no mask annotations, then disable loss computing for them
            for idx, e in enumerate(roidb):
                has_mask = e['has_mask']
                if has_mask:
                    continue
                ind = rpn_ret['mask_rois'][:, 0] == idx
                rpn_ret['masks_int32'][ind, :] = np.zeros_like(
                    rpn_ret['masks_int32'][ind, :])

            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
示例#6
0
    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 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 for training
            if not self.training:
                return_dict['cls_score'] = rpn_ret['rpn_roi_probs']
                return_dict['bbox_pred'] = rpn_ret['rpn_rois']
            else:
                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
            if not cfg.MODEL.RPN_ONLY:
                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

            # mask loss
            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']
            if not cfg.MODEL.RPN_ONLY:
                return_dict['cls_score'] = cls_score
                return_dict['bbox_pred'] = bbox_pred
            if cfg.PRM_ON:
                return_dict['rpn_ret'] = rpn_ret

        return return_dict
    def _forward(self,
                 data,
                 im_info,
                 dataset_name=None,
                 roidb=None,
                 use_gt_labels=False,
                 include_feat=False,
                 **rpn_kwargs):
        im_data = data
        if self.training:
            roidb = list(map(lambda x: blob_utils.deserialize(x)[0], roidb))
        if dataset_name is not None:
            dataset_name = blob_utils.deserialize(dataset_name)
        else:
            dataset_name = cfg.TRAIN.DATASETS[
                0] if self.training else cfg.TEST.DATASETS[
                    0]  # assuming only one dataset per run

        device_id = im_data.get_device()

        return_dict = {}  # A dict to collect return variables

        blob_conv = self.Conv_Body(im_data)
        blob_conv_prd = self.Prd_RCNN.Conv_Body(im_data)

        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:]
            blob_conv_prd = blob_conv_prd[-self.num_roi_levels:]

        if not cfg.TRAIN.USE_GT_BOXES:
            rpn_ret = self.RPN(blob_conv, im_info, roidb)

            if cfg.MODEL.SHARE_RES5 and self.training:
                box_feat, res5_feat = self.Box_Head(blob_conv,
                                                    rpn_ret,
                                                    use_relu=True)
            else:
                box_feat = self.Box_Head(blob_conv, rpn_ret, use_relu=True)
            cls_score, bbox_pred = self.Box_Outs(box_feat)

        # now go through the predicate branch
        use_relu = False if cfg.MODEL.NO_FC7_RELU else True
        if self.training:
            if cfg.TRAIN.USE_GT_BOXES:
                # we always feed one image per batch during training
                assert len(roidb) == 1
                im_scale = im_info.data.numpy()[:, 2][0]
                im_w = im_info.data.numpy()[:, 1][0]
                im_h = im_info.data.numpy()[:, 0][0]
                sbj_boxes = roidb[0]['sbj_gt_boxes']
                obj_boxes = roidb[0]['obj_gt_boxes']
                sbj_all_boxes = _augment_gt_boxes_by_perturbation(
                    sbj_boxes, im_w, im_h)
                obj_all_boxes = _augment_gt_boxes_by_perturbation(
                    obj_boxes, im_w, im_h)
                det_all_boxes = np.vstack((sbj_all_boxes, obj_all_boxes))
                det_all_boxes = np.unique(det_all_boxes, axis=0)
                det_all_rois = det_all_boxes * im_scale
                repeated_batch_idx = 0 * blob_utils.ones(
                    (det_all_rois.shape[0], 1))
                det_all_rois = np.hstack((repeated_batch_idx, det_all_rois))
                rel_ret = self.RelPN(det_all_rois, None, None, im_info,
                                     dataset_name, roidb)
            else:
                fg_inds = np.where(rpn_ret['labels_int32'] > 0)[0]
                det_rois = rpn_ret['rois'][fg_inds]
                det_labels = rpn_ret['labels_int32'][fg_inds]
                det_scores = F.softmax(cls_score[fg_inds], dim=1)
                rel_ret = self.RelPN(det_rois, det_labels, det_scores, im_info,
                                     dataset_name, roidb)
            sbj_feat = self.Box_Head(blob_conv,
                                     rel_ret,
                                     rois_name='sbj_rois',
                                     use_relu=use_relu)
            obj_feat = self.Box_Head(blob_conv,
                                     rel_ret,
                                     rois_name='obj_rois',
                                     use_relu=use_relu)
        else:
            if roidb is not None:
                im_scale = im_info.data.numpy()[:, 2][0]
                im_w = im_info.data.numpy()[:, 1][0]
                im_h = im_info.data.numpy()[:, 0][0]
                sbj_boxes = roidb['sbj_gt_boxes']
                obj_boxes = roidb['obj_gt_boxes']
                sbj_rois = sbj_boxes * im_scale
                obj_rois = obj_boxes * im_scale
                repeated_batch_idx = 0 * blob_utils.ones(
                    (sbj_rois.shape[0], 1))
                sbj_rois = np.hstack((repeated_batch_idx, sbj_rois))
                obj_rois = np.hstack((repeated_batch_idx, obj_rois))
                rel_rois = box_utils.rois_union(sbj_rois, obj_rois)
                rel_ret = {}
                rel_ret['sbj_rois'] = sbj_rois
                rel_ret['obj_rois'] = obj_rois
                rel_ret['rel_rois'] = rel_rois
                if cfg.FPN.FPN_ON and cfg.FPN.MULTILEVEL_ROIS:
                    lvl_min = cfg.FPN.ROI_MIN_LEVEL
                    lvl_max = cfg.FPN.ROI_MAX_LEVEL
                    rois_blob_names = ['sbj_rois', 'obj_rois', 'rel_rois']
                    for rois_blob_name in rois_blob_names:
                        # Add per FPN level roi blobs named like: <rois_blob_name>_fpn<lvl>
                        target_lvls = fpn_utils.map_rois_to_fpn_levels(
                            rel_ret[rois_blob_name][:, 1:5], lvl_min, lvl_max)
                        fpn_utils.add_multilevel_roi_blobs(
                            rel_ret, rois_blob_name, rel_ret[rois_blob_name],
                            target_lvls, lvl_min, lvl_max)
                if use_gt_labels:
                    sbj_labels = roidb['sbj_gt_classes']  # start from 0
                    obj_labels = roidb['obj_gt_classes']  # start from 0
                    sbj_scores = np.ones_like(sbj_labels, dtype=np.float32)
                    obj_scores = np.ones_like(obj_labels, dtype=np.float32)
                else:
                    sbj_det_feat = self.Box_Head(blob_conv,
                                                 rel_ret,
                                                 rois_name='sbj_rois',
                                                 use_relu=True)
                    sbj_cls_scores, _ = self.Box_Outs(sbj_det_feat)
                    sbj_cls_scores = sbj_cls_scores.data.cpu().numpy()
                    obj_det_feat = self.Box_Head(blob_conv,
                                                 rel_ret,
                                                 rois_name='obj_rois',
                                                 use_relu=True)
                    obj_cls_scores, _ = self.Box_Outs(obj_det_feat)
                    obj_cls_scores = obj_cls_scores.data.cpu().numpy()
                    sbj_labels = np.argmax(sbj_cls_scores[:, 1:], axis=1)
                    obj_labels = np.argmax(obj_cls_scores[:, 1:], axis=1)
                    sbj_scores = np.amax(sbj_cls_scores[:, 1:], axis=1)
                    obj_scores = np.amax(obj_cls_scores[:, 1:], axis=1)
                rel_ret['sbj_scores'] = sbj_scores.astype(np.float32,
                                                          copy=False)
                rel_ret['obj_scores'] = obj_scores.astype(np.float32,
                                                          copy=False)
                rel_ret['sbj_labels'] = sbj_labels.astype(
                    np.int32, copy=False) + 1  # need to start from 1
                rel_ret['obj_labels'] = obj_labels.astype(
                    np.int32, copy=False) + 1  # need to start from 1
                rel_ret['all_sbj_labels_int32'] = sbj_labels.astype(np.int32,
                                                                    copy=False)
                rel_ret['all_obj_labels_int32'] = obj_labels.astype(np.int32,
                                                                    copy=False)
                sbj_feat = self.Box_Head(blob_conv,
                                         rel_ret,
                                         rois_name='sbj_rois',
                                         use_relu=use_relu)
                obj_feat = self.Box_Head(blob_conv,
                                         rel_ret,
                                         rois_name='obj_rois',
                                         use_relu=use_relu)
            else:
                score_thresh = cfg.TEST.SCORE_THRESH
                while score_thresh >= -1e-06:  # a negative value very close to 0.0
                    det_rois, det_labels, det_scores = \
                        self.prepare_det_rois(rpn_ret['rois'], cls_score, bbox_pred, im_info, score_thresh)
                    rel_ret = self.RelPN(det_rois, det_labels, det_scores,
                                         im_info, dataset_name, roidb)
                    valid_len = len(rel_ret['rel_rois'])
                    if valid_len > 0:
                        break
                    logger.info(
                        'Got {} rel_rois when score_thresh={}, changing to {}'.
                        format(valid_len, score_thresh, score_thresh - 0.01))
                    score_thresh -= 0.01
                det_feat = self.Box_Head(blob_conv,
                                         rel_ret,
                                         rois_name='det_rois',
                                         use_relu=use_relu)
                sbj_feat = det_feat[rel_ret['sbj_inds']]
                obj_feat = det_feat[rel_ret['obj_inds']]

        rel_feat = self.Prd_RCNN.Box_Head(blob_conv_prd,
                                          rel_ret,
                                          rois_name='rel_rois',
                                          use_relu=use_relu)

        concat_feat = torch.cat((sbj_feat, rel_feat, obj_feat), dim=1)

        if cfg.MODEL.USE_FREQ_BIAS or cfg.MODEL.RUN_BASELINE or cfg.MODEL.USE_SEM_CONCAT:
            sbj_labels = rel_ret['all_sbj_labels_int32']
            obj_labels = rel_ret['all_obj_labels_int32']
        else:
            sbj_labels = None
            obj_labels = None

        # when MODEL.USE_SEM_CONCAT, memory runs out if the whole batch is fed once
        # so we need to feed the batch twice if it's big
        gn_size = 1000
        if cfg.MODEL.USE_SEM_CONCAT and concat_feat.shape[0] > gn_size:
            group = int(math.ceil(concat_feat.shape[0] / gn_size))
            prd_cls_scores = None
            sbj_cls_scores = None
            obj_cls_scores = None
            for i in range(group):
                end = int(min((i + 1) * gn_size, concat_feat.shape[0]))
                concat_feat_i = concat_feat[i * gn_size:end]
                sbj_labels_i = sbj_labels[
                    i * gn_size:end] if sbj_labels is not None else None
                obj_labels_i = obj_labels[
                    i * gn_size:end] if obj_labels is not None else None
                sbj_feat_i = sbj_feat[i * gn_size:end]
                obj_feat_i = obj_feat[i * gn_size:end]
                prd_cls_scores_i, sbj_cls_scores_i, obj_cls_scores_i = \
                    self.RelDN(concat_feat_i, sbj_labels_i, obj_labels_i, sbj_feat_i, obj_feat_i)
                if prd_cls_scores is None:
                    prd_cls_scores = prd_cls_scores_i
                    sbj_cls_scores = sbj_cls_scores_i
                    obj_cls_scores = obj_cls_scores_i
                else:
                    prd_cls_scores = torch.cat(
                        (prd_cls_scores, prd_cls_scores_i))
                    sbj_cls_scores = torch.cat(
                        (sbj_cls_scores, sbj_cls_scores_i
                         )) if sbj_cls_scores_i is not None else sbj_cls_scores
                    obj_cls_scores = torch.cat(
                        (obj_cls_scores, obj_cls_scores_i
                         )) if obj_cls_scores_i is not None else obj_cls_scores
        else:
            prd_cls_scores, sbj_cls_scores, obj_cls_scores = \
                    self.RelDN(concat_feat, sbj_labels, obj_labels, sbj_feat, obj_feat)

        if self.training:
            return_dict['losses'] = {}
            return_dict['metrics'] = {}
            if not cfg.TRAIN.USE_GT_BOXES:
                # 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
            loss_cls_prd, accuracy_cls_prd = reldn_heads.reldn_losses(
                prd_cls_scores,
                rel_ret['all_prd_labels_int32'],
                weight=self.prd_weights)
            return_dict['losses']['loss_cls_prd'] = loss_cls_prd
            return_dict['metrics']['accuracy_cls_prd'] = accuracy_cls_prd
            if cfg.MODEL.USE_SEPARATE_SO_SCORES:
                loss_cls_sbj, accuracy_cls_sbj = reldn_heads.reldn_losses(
                    sbj_cls_scores,
                    rel_ret['all_sbj_labels_int32'],
                    weight=self.obj_weights)
                return_dict['losses']['loss_cls_sbj'] = loss_cls_sbj
                return_dict['metrics']['accuracy_cls_sbj'] = accuracy_cls_sbj
                loss_cls_obj, accuracy_cls_obj = reldn_heads.reldn_losses(
                    obj_cls_scores,
                    rel_ret['all_obj_labels_int32'],
                    weight=self.obj_weights)
                return_dict['losses']['loss_cls_obj'] = loss_cls_obj
                return_dict['metrics']['accuracy_cls_obj'] = accuracy_cls_obj

            if cfg.TRAIN.HUBNESS:
                loss_hubness_prd = reldn_heads.add_hubness_loss(prd_cls_scores)
                loss_hubness_sbj = reldn_heads.add_hubness_loss(sbj_cls_scores)
                loss_hubness_obj = reldn_heads.add_hubness_loss(obj_cls_scores)
                return_dict['losses']['loss_hubness_prd'] = loss_hubness_prd
                return_dict['losses']['loss_hubness_sbj'] = loss_hubness_sbj
                return_dict['losses']['loss_hubness_obj'] = loss_hubness_obj

            # 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['sbj_rois'] = rel_ret['sbj_rois']
            return_dict['obj_rois'] = rel_ret['obj_rois']
            return_dict['sbj_labels'] = rel_ret['sbj_labels']
            return_dict['obj_labels'] = rel_ret['obj_labels']
            return_dict['sbj_scores'] = rel_ret['sbj_scores']
            return_dict['sbj_scores_out'] = sbj_cls_scores
            return_dict['obj_scores'] = rel_ret['obj_scores']
            return_dict['obj_scores_out'] = obj_cls_scores
            return_dict['prd_scores'] = prd_cls_scores
            if include_feat:
                return_dict['sbj_feat'] = sbj_feat
                return_dict['obj_feat'] = obj_feat
                return_dict['prd_feat'] = concat_feat

        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)
        '''
        anchors of different scale and aspect ratios will be collected
        the function self.RPN is rpn_heads.generic_rpn_outputs
        it will call FPN.fpn_rpn_outputs which is a class name, so it creates the class object
        '''
        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:
                '''in faster rcnn fpn r-50 case it is fast_rcnn_heads.roi_2mlp_head'''
                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, support_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
        original_blob_conv = self.Conv_Body(im_data)

        if not self.training:
            all_cls = roidb[0]['support_cls']
            test_way = len(all_cls)
            test_shot = int(roidb[0]['support_boxes'].shape[0] / test_way)
        else:
            train_way = 2  #2
            train_shot = 5  #5
        support_blob_conv = self.Conv_Body(support_data.squeeze(0))

        img_num = int(support_blob_conv.shape[0])
        img_channel = int(original_blob_conv.shape[1])
        # Construct support rpn_ret
        support_rpn_ret = {
            'rois': np.insert(roidb[0]['support_boxes'][0], 0,
                              0.)[np.newaxis, :]
        }
        if img_num > 1:
            for i in range(img_num - 1):
                support_rpn_ret['rois'] = np.concatenate(
                    (support_rpn_ret['rois'],
                     np.insert(roidb[0]['support_boxes'][i + 1], 0,
                               float(i + 1))[np.newaxis, :]),
                    axis=0)
        # Get support pooled feature
        support_feature = self.roi_feature_transform(
            support_blob_conv,
            support_rpn_ret,
            blob_rois='rois',
            method=cfg.FAST_RCNN.ROI_XFORM_METHOD,
            resolution=cfg.FAST_RCNN.ROI_XFORM_RESOLUTION,
            spatial_scale=self.Conv_Body.spatial_scale,
            sampling_ratio=cfg.FAST_RCNN.ROI_XFORM_SAMPLING_RATIO)

        blob_conv = original_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

        assert not cfg.MODEL.RPN_ONLY

        support_box_feat = self.Box_Head(support_blob_conv, support_rpn_ret)
        if self.training:
            support_feature_mean_0 = support_feature[:5].mean(0, True)
            support_pool_feature_0 = self.avgpool(support_feature_mean_0)
            correlation_0 = F.conv2d(original_blob_conv,
                                     support_pool_feature_0.permute(
                                         1, 0, 2, 3),
                                     groups=1024)
            rpn_ret_0 = self.RPN(correlation_0, im_info, roidb)
            box_feat_0 = self.Box_Head(blob_conv, rpn_ret_0)
            support_0 = support_box_feat[:5].mean(
                0, True)  # simple average few shot support features

            support_feature_mean_1 = support_feature[5:10].mean(0, True)
            support_pool_feature_1 = self.avgpool(support_feature_mean_1)
            correlation_1 = F.conv2d(original_blob_conv,
                                     support_pool_feature_1.permute(
                                         1, 0, 2, 3),
                                     groups=1024)
            rpn_ret_1 = self.RPN(correlation_1, im_info, roidb)
            box_feat_1 = self.Box_Head(blob_conv, rpn_ret_1)
            support_1 = support_box_feat[5:10].mean(
                0, True)  # simple average few shot support features

            cls_score_now_0, bbox_pred_now_0 = self.Box_Outs(
                box_feat_0, support_0)
            cls_score_now_1, bbox_pred_now_1 = self.Box_Outs(
                box_feat_1, support_1)

            cls_score = torch.cat([cls_score_now_0, cls_score_now_1], dim=0)
            bbox_pred = torch.cat([bbox_pred_now_0, bbox_pred_now_1], dim=0)
            rpn_ret = {}
            for key in rpn_ret_0.keys():
                if key != 'rpn_cls_logits' and key != 'rpn_bbox_pred':
                    rpn_ret[key] = rpn_ret_0[
                        key]  #np.concatenate((rpn_ret_0[key], rpn_ret_1[key]), axis=0)
                else:
                    rpn_ret[key] = torch.cat([rpn_ret_0[key], rpn_ret_1[key]],
                                             dim=0)

        else:
            for way_id in range(test_way):
                begin = way_id * test_shot
                end = (way_id + 1) * test_shot

                support_feature_mean = support_feature[begin:end].mean(0, True)
                support_pool_feature = self.avgpool(support_feature_mean)
                correlation = F.conv2d(original_blob_conv,
                                       support_pool_feature.permute(
                                           1, 0, 2, 3),
                                       groups=1024)
                rpn_ret = self.RPN(correlation, im_info, roidb)

                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)

                    support = support_box_feat[begin:end].mean(
                        0, True)  # simple average few shot support features

                    cls_score_now, bbox_pred_now = self.Box_Outs(
                        box_feat, support)
                    cls_now = cls_score_now.new_full(
                        (cls_score_now.shape[0], 1),
                        int(roidb[0]['support_cls'][way_id]))
                    cls_score_now = torch.cat([cls_score_now, cls_now], dim=1)
                    rois = rpn_ret['rois']
                    if way_id == 0:
                        cls_score = cls_score_now
                        bbox_pred = bbox_pred_now
                        rois_all = rois
                    else:
                        cls_score = torch.cat([cls_score, cls_score_now],
                                              dim=0)
                        bbox_pred = torch.cat([bbox_pred, bbox_pred_now],
                                              dim=0)
                        rois_all = np.concatenate((rois_all, rois), axis=0)
                else:
                    # TODO: complete the returns for RPN only situation
                    pass
            rpn_ret['rois'] = rois_all
        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'))))
            target_cls = roidb[0]['target_cls']
            rpn_ret['labels_obj'] = np.array(
                [int(i == target_cls) for i in rpn_ret['labels_int32']])

            # filter other class bbox targets, only supervise the target cls bbox, because the bbox_pred is from the compare feature.
            bg_idx = np.where(rpn_ret['labels_int32'] != target_cls)[0]
            rpn_ret['bbox_targets'][bg_idx] = np.full_like(
                rpn_ret['bbox_targets'][bg_idx], 0.)
            rpn_ret['bbox_inside_weights'][bg_idx] = np.full_like(
                rpn_ret['bbox_inside_weights'][bg_idx], 0.)
            rpn_ret['bbox_outside_weights'][bg_idx] = np.full_like(
                rpn_ret['bbox_outside_weights'][bg_idx], 0.)

            neg_labels_obj = np.full_like(rpn_ret['labels_obj'], 0.)
            neg_bbox_targets = np.full_like(rpn_ret['bbox_targets'], 0.)
            neg_bbox_inside_weights = np.full_like(
                rpn_ret['bbox_inside_weights'], 0.)
            neg_bbox_outside_weights = np.full_like(
                rpn_ret['bbox_outside_weights'], 0.)

            rpn_ret['labels_obj'] = np.concatenate(
                [rpn_ret['labels_obj'], neg_labels_obj], axis=0)
            rpn_ret['bbox_targets'] = np.concatenate(
                [rpn_ret['bbox_targets'], neg_bbox_targets], axis=0)
            rpn_ret['bbox_inside_weights'] = np.concatenate(
                [rpn_ret['bbox_inside_weights'], neg_bbox_inside_weights],
                axis=0)
            rpn_ret['bbox_outside_weights'] = np.concatenate(
                [rpn_ret['bbox_outside_weights'], neg_bbox_outside_weights],
                axis=0)

            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_obj'],
                rpn_ret['bbox_targets'], rpn_ret['bbox_inside_weights'],
                rpn_ret['bbox_outside_weights'])
            return_dict['losses']['loss_cls'] = 1 * loss_cls
            return_dict['losses']['loss_bbox'] = 1 * 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
示例#10
0
    def _forward(self,
                 data,
                 im_info,
                 roidb=None,
                 data_flow=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)
        assert (len(blob_conv) == 5)

        # if gru is for dynamic model, we train gru only.
        #if cfg.CONVGRU.DYNAMIC_MODEL:
        #   for i in range(5):
        #      blob_conv[i].detach_()

        # hidden states management.
        if not cfg.CONVGRU.DYNAMIC_MODEL:
            #Every time, the image size may be different, so we create new hidden states.
            self.create_hidden_states(blob_conv)
        else:
            #Caller needs to manually clean the hidden states.
            self.check_exist_hidden_states(blob_conv)

        if cfg.MODEL.USE_DELTA_FLOW:
            #overwrite dataflow.
            data_flow = torch.zeros(im_data.shape[0],
                                    2,
                                    im_data.shape[2],
                                    im_data.shape[3],
                                    device=im_data.device)
            for i in range(5):
                delta_flow_feature = self.Conv_Delta_Flows_Features[i](
                    blob_conv[i].detach())
                delta_flow_lvl = None
                if not self.flow_features[i] is None:
                    delta_flow_lvl = delta_flow_feature - self.flow_features[i]
                self.flow_features[i] = delta_flow_feature
                if not delta_flow_lvl is None:
                    delta_flow_lvl = nn.functional.upsample(delta_flow_lvl,
                                                            scale_factor=1.0 /
                                                            self.fpn_scales[i],
                                                            mode='bilinear')
                    delta_flow_lvl = delta_flow_lvl / self.fpn_scales[i]
                    data_flow = data_flow + delta_flow_lvl
        if cfg.CONVGRU.DYNAMIC_MODEL is True:
            #if dynamic model, hidden_states need to be updated.
            warped_hidden_states = [None] * 5
            for i in range(5):
                if not data_flow is None:
                    #assert not np.any(np.isnan(data_flow.data.cpu().numpy())), 'data_flow has nan.'
                    warped_hidden_states[i] = self.FlowAligns[i](
                        self.hidden_states[i], data_flow)
                else:
                    warped_hidden_states[i] = self.hidden_states[i]

        for i in range(4, -1, -1):
            if cfg.CONVGRU.DYNAMIC_MODEL is True:
                blob_conv[i] = self.ConvGRUs[i](
                    (blob_conv[i], warped_hidden_states[i]))
            else:
                blob_conv[i] = self.ConvGRUs[i](
                    (blob_conv[i], self.hidden_states[i]))
            if i < 4:
                blob_conv[i] = blob_conv[i] / 2.0 + nn.functional.upsample(
                    blob_conv[i + 1], scale_factor=0.5, mode='bilinear') / 2.0

            if self.update_hidden_states and cfg.CONVGRU.DYNAMIC_MODEL:
                self.hidden_states[i] = blob_conv[i]

        if self.stop_after_hidden_states is True:
            # print('stop after hidden states.')
            return None

        rpn_ret = self.RPN(blob_conv, im_info, roidb)

        # Handle blob with no GT boxes.
        HEAD_TRAIN = True
        if cfg.CONVGRU.DYNAMIC_MODEL and not roidb is None:
            assert (len(roidb) == 1)
            if not len(rpn_ret['bbox_targets']) > 0:
                HEAD_TRAIN = False
        # 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 HEAD_TRAIN:
                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)
                if not cfg.MODEL.IDENTITY_TRAINING:
                    cls_score, bbox_pred = self.Box_Outs(box_feat)
                else:
                    cls_score, bbox_pred, id_score = 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
            if cfg.MODEL.ADD_UNKNOWN_CLASS is True:
                cls_weights = torch.tensor([1.] * (cfg.MODEL.NUM_CLASSES - 1) +
                                           [0.],
                                           device=cls_score.device)
            else:
                cls_weights = None

            if HEAD_TRAIN:
                if not cfg.MODEL.IDENTITY_TRAINING:
                    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'],
                        cls_weights=cls_weights)
                else:
                    loss_cls, loss_bbox, accuracy_cls, loss_id, accuracy_id = 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'],
                        cls_weights=cls_weights,
                        id_score=id_score,
                        id_int32=rpn_ret['global_id_int32'])
                    return_dict['losses']['loss_id'] = loss_id
                    return_dict['metrics']['accuracy_id'] = accuracy_id
                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
            if cfg.MODEL.IDENTITY_TRAINING:
                return_dict['id_score'] = id_score

        return return_dict
示例#11
0
    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 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)
            elif cfg.MODEL.FASTER_RCNN: #hw FAST_RCNN or LIGHT_HEAD_RCNN网络预测类别和坐标
                box_feat = self.Box_Head(blob_conv, rpn_ret)
                cls_score, bbox_pred = self.Box_Outs(box_feat)
                return_dict['cls_score'] = cls_score
                return_dict['bbox_pred'] = bbox_pred
            elif cfg.MODEL.LIGHT_HEAD_RCNN:
                box_feat = F.relu(self.LightHead(blob_conv), inplace=True)
                cls_score, bbox_pred = self.Box_Outs(box_feat,rpn_ret)
                return_dict['cls_score'] = cls_score
                return_dict['bbox_pred'] = bbox_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'))
            )) #hw RPN网络类别和坐标的损失值计算
            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 #hw RCNN网络类别和坐标的损失值计算
            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
            elif cfg.MODEL.LIGHT_HEAD_RCNN:  # hw
                loss_cls, loss_bbox = light_head_rcnn_heads.light_head_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



            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
示例#12
0
    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
示例#14
0
    def _forward(self, blob_conv, rpn_ret, adv_target=None, flags=None):
        return_dict = {}

        batch_size = blob_conv.size(0)
        if self.training and cfg.DEBUG:
            # debug: batch_size and fg-fraction
            fg = len([x for x in rpn_ret['labels_int32'] if x > 0])
            print("\tbatch-size in discriminator: {} (fg: {}%)".format(
                batch_size, 1.0 * fg / batch_size * 100.0))

            print("\tBlob_conv size in discriminator: {}".format(
                blob_conv.view(batch_size, -1).size()))

        blob_conv_flattened = blob_conv.view(batch_size, -1)

        adv_score = self.adversarial(blob_conv_flattened)

        box_feat = self.Box_Head(blob_conv_flattened)
        cls_score, bbox_pred = self.Box_Outs(box_feat)

        if self.training:
            if adv_target is None:
                raise ValueError(
                    "adv_target must not be None during training!!")

            return_dict['losses'] = {}
            return_dict['metrics'] = {}

            if cfg.GAN.TRAIN.IGNORE_BG_ADV_LOSS:
                # ignore adversarial loss for background RoIs
                mask = np.where(rpn_ret['labels_int32'] == 0)
                fg = len([x for x in rpn_ret['labels_int32'] if x > 0])
                bg = len([x for x in rpn_ret['labels_int32'] if x == 0])
                if cfg.DEBUG:
                    print(
                        "ignoring backgound rois in adv_loss: {} / {}".format(
                            bg, len(rpn_ret['labels_int32'])))

                loss_adv = self.adversarial_loss(adv_score,
                                                 adv_target,
                                                 reduce=False)
                loss_adv[mask] = 0.0
                if fg > 0:
                    loss_adv = loss_adv * len(rpn_ret['labels_int32']) / fg
                loss_adv = loss_adv.mean()
            else:
                loss_adv = self.adversarial_loss(adv_score, adv_target)

            return_dict['losses']['loss_adv'] = loss_adv

            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

            # 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 testing
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_pred
            return_dict['adv_score'] = adv_score
            return_dict['rois'] = rpn_ret['rois']
            return_dict['rpn_ret'] = rpn_ret

        return return_dict
示例#15
0
    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.RPN.VIS_QUANT_TARGET:
                import numpy as np
                import json
                import os
                import time

                gt_boxes = []
                gt_label = roidb[0]['gt_classes']
                for inds, item in enumerate(gt_label):
                    if item != 0:
                        gt_boxes.append(roidb[0]['boxes'][inds])

                gt_boxes = np.array(gt_boxes, dtype=np.float32)
                gt_boxes *= im_info.detach().numpy()[:, 2]

                path = "/nfs/project/libo_i/Boosting/Targets_Info"
                if not os.path.exists(path):
                    os.makedirs(path)
                b, c, h, w = rpn_kwargs['rpn_cls_logits_fpn3'].shape
                sample_targets = rpn_kwargs[
                    'rpn_bbox_targets_wide_fpn3'][:, :, :h, :w]

                line_targets = sample_targets.detach().data.cpu().numpy()
                with open(os.path.join(path, "quant_anchors.json"), "r") as fp:
                    quant_anchors = np.array(json.load(fp), dtype=np.float32)
                    quant_anchors = quant_anchors[:h, :w]

                line_targets = line_targets[:, 4:8, :, :].transpose(
                    (0, 2, 3, 1)).reshape(quant_anchors.shape)
                line_targets = line_targets.reshape(-1, 4)

                width = im_data.shape[3]
                height = im_data.shape[2]
                # 在这里加上targets的偏移

                line_quant_anchors = quant_anchors.reshape(-1, 4)
                pred_boxes = box_utils.onedim_bbox_transform(
                    line_quant_anchors, line_targets)
                pred_boxes = box_utils.clip_tiled_boxes(
                    pred_boxes, (height, width, 3))

                im = im_data.detach().cpu().numpy().reshape(3, height,
                                                            width).transpose(
                                                                (1, 2, 0))

                means = np.squeeze(cfg.PIXEL_MEANS)
                for i in range(3):
                    im[:, :, i] += means[i]

                im = im.astype(int)
                dpi = 200
                fig = plt.figure(frameon=False)
                fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
                ax = plt.Axes(fig, [0., 0., 1., 1.])
                ax.axis('off')
                fig.add_axes(ax)
                ax.imshow(im[:, :, ::-1])
                # 在im上添加gt
                for item in gt_boxes:
                    ax.add_patch(
                        plt.Rectangle((item[0], item[1]),
                                      item[2] - item[0],
                                      item[3] - item[1],
                                      fill=False,
                                      edgecolor='white',
                                      linewidth=1,
                                      alpha=1))

                cnt = 0
                for inds, before_item in enumerate(line_quant_anchors):
                    after_item = pred_boxes[inds]
                    targets_i = line_targets[inds]
                    if np.sum(targets_i) == 0:
                        continue
                    ax.add_patch(
                        plt.Rectangle((before_item[0], before_item[1]),
                                      before_item[2] - before_item[0],
                                      before_item[3] - before_item[1],
                                      fill=False,
                                      edgecolor='r',
                                      linewidth=1,
                                      alpha=1))

                    ax.add_patch(
                        plt.Rectangle((after_item[0], after_item[1]),
                                      after_item[2] - after_item[0],
                                      after_item[3] - after_item[1],
                                      fill=False,
                                      edgecolor='g',
                                      linewidth=1,
                                      alpha=1))

                    logger.info("valid boxes: {}".format(cnt))
                    cnt += 1

                if cnt != 0:
                    ticks = time.time()
                    fig.savefig(
                        "/nfs/project/libo_i/Boosting/Targets_Info/{}.png".
                        format(ticks),
                        dpi=dpi)

                plt.close('all')

            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'])
            if cfg.RPN.ZEROLOSS:
                zero_loss_bbox = torch.Tensor([0.]).squeeze().cuda()
                zero_loss_bbox.requires_grad = True
                zero_loss_cls = torch.Tensor([0.]).squeeze().cuda()
                zero_loss_cls.requires_grad = True
                return_dict['losses']['loss_bbox'] = zero_loss_bbox
                return_dict['losses']['loss_cls'] = zero_loss_cls

            else:
                return_dict['losses']['loss_bbox'] = loss_bbox
                return_dict['losses']['loss_cls'] = loss_cls

            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
            if cfg.TEST.PROPOSALS_OUT:
                import os
                import json
                import numpy as np

                # 直接在这里做变换,输出经过变换之后的1000个框
                bbox_pred = bbox_pred.data.cpu().numpy().squeeze()
                box_deltas = bbox_pred.reshape([-1, bbox_pred.shape[-1]])
                shift_boxes = box_utils.bbox_transform(
                    rpn_ret['rois'][:, 1:5], 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 = []
                inds_all = []
                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]
                    onecls_pred_boxes += boxes_j.tolist()
                    inds_all.extend(inds.tolist())

                inds_all = np.array(inds_all, dtype=np.int)
                aligned_proposals = rpn_ret['rois'][:, 1:5][inds_all]
                aligned_boxes = np.array(onecls_pred_boxes, dtype=np.float32)

                assert inds_all.shape[0] == aligned_boxes.shape[0]
                assert aligned_proposals.size == aligned_boxes.size

                path = "/nfs/project/libo_i/Boosting/Anchor_Info"
                with open(os.path.join(path, "proposals.json"), "w") as fp:
                    json.dump(aligned_proposals.tolist(), fp)

                with open(os.path.join(path, "boxes.json"), "w") as fp:
                    json.dump(aligned_boxes.tolist(), fp)

        return return_dict
示例#16
0
    def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
        M = cfg.LESION.NUM_IMAGES_3DCE
        # data.shape: [n,M*3,h,w]
        n, c, h, w = data.shape
        im_data = data.view(n * M, 3, h, w)
        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
        if cfg.LESION.USE_POSITION:
            blob_conv, res5_feat = self.Conv_Body(im_data)
        else:
            blob_conv = self.Conv_Body(im_data)

        # blob.shape = [n, c,h,w] for 2d
        # blob.shape = [nM,c,h,w] for 3DCE or multi-modal

        blob_conv_for_RPN = []
        blob_conv_for_RCNN = []
        # Used for MVP-Net, concat all slices before RPN.
        if cfg.LESION.CONCAT_BEFORE_RPN:
            for blob in blob_conv:
                _, c, h, w = blob.shape
                blob = blob.view(n, M * c, h, w)
                blob_cbam = self.cbam(blob)
                blob_conv_for_RPN.append(blob_cbam)
            blob_conv_for_RCNN = blob_conv_for_RPN
        # Used for Standard 3DCE, feed middle slice into RPN.
        else:
            for blob in blob_conv:
                _, c, h, w = blob.shape
                blob = blob.view(n, M * c, h, w)
                blob_conv_for_RPN.append(blob[:, (M // 2) * c:(M // 2 + 1) *
                                              c, :, :])
                blob_conv_for_RCNN.append(blob)

        rpn_ret = self.RPN(blob_conv_for_RPN, 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_for_RCNN = blob_conv_for_RCNN[-self.num_roi_levels:]

        if not self.training:
            return_dict['blob_conv'] = blob_conv_for_RCNN

        if not cfg.MODEL.RPN_ONLY:
            if cfg.MODEL.SHARE_RES5 and self.training:
                box_feat, res5_feat = self.Box_Head(blob_conv_for_RCNN,
                                                    rpn_ret)
            else:
                box_feat = self.Box_Head(blob_conv_for_RCNN, rpn_ret)

            if cfg.LESION.USE_POSITION:
                if cfg.LESION.NUM_IMAGES_3DCE == 3:
                    position_feat = blob_conv_for_RPN[0]
                    n, c, h, w = position_feat.shape
                    position_feat = position_feat.view(3, 256, h, w)
                    #print(position_feat.shape)
                elif cfg.LESION.NUM_IMAGES_3DCE == 9:
                    position_feat = blob_conv_for_RPN[0][:,
                                                         3 * 256:6 * 256, :, :]
                    n, c, h, w = position_feat.shape
                    position_feat = position_feat.view(3, 256, h, w)
                    #print(position_feat.shape)
                position_feat = self.Position_Head(position_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
            if cfg.LESION.POS_CONCAT_RCNN:
                try:
                    if self.training:
                        # expand position_feat as box_feat for concatenation.
                        # box_feat: (n*num_nms, 1024)
                        # position_feat: (n, 1024) expand--> position_feat_:(num_nms, n ,1024)
                        position_feat_ = position_feat.expand(
                            cfg.TRAIN.RPN_BATCH_SIZE_PER_IM, -1, -1)
                        # position_feat_: (num_nms ,n , 1024)  transpose--> (n, num_nms, 1024)  view-->(n*num_nms, 1024)
                        position_feat_concat = torch.transpose(
                            position_feat_, 1, 0).contiguous().view(-1, 1024)
                    else:
                        # box_feat: (1, 1024)
                        # position_feat: (1, 1024), position_feat_:(1000 ,1024)
                        position_feat_ = position_feat.expand(
                            cfg.TEST.RPN_PRE_NMS_TOP_N, -1)
                        # position_feat_: (1, 1024)  transpose--> position_feat_concat (n, num_nms, 1024)
                        position_feat_concat = torch.transpose(
                            position_feat_, 1,
                            0).contiguous().view_as(box_feat)
                except:
                    import pdb
                    pdb.set_trace()
                    print('position_feat', position_feat.shape)
                    print('box_feat', box_feat.shape)
                    print('position_feat_', position_feat_.shape)
                    print('position_feat_concat', position_feat_concat.shape)
                else:
                    box_feat = torch.cat([box_feat, position_feat_concat], 1)
                    cls_score, bbox_pred = self.Box_Outs(box_feat)
            cls_score, bbox_pred = self.Box_Outs(box_feat)
            # print cls_score.shape
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_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.LESION.USE_POSITION:
                pos_cls_loss, pos_reg_loss, accuracy_position = position_losses(
                    pos_cls_pred, pos_reg_pred, roidb)
                return_dict['losses']['pos_cls_loss'] = pos_cls_loss
                return_dict['losses']['pos_reg_loss'] = pos_reg_loss
                return_dict['metrics']['accuracy_position'] = accuracy_position

            # 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
示例#17
0
    def forward(self, data, im_info=None, roidb=None, only_body=None, **rpn_kwargs):
        im_data = data
        if self.training and cfg.RPN.RPN_ON:
            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)

        if cfg.RPN.RPN_ON:
            rpn_ret = self.RPN(blob_conv, im_info, roidb)
        else:
            rpn_ret = rpn_kwargs

        # 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 only_body is not None:
            return return_dict

        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 cfg.RPN.RPN_ON:
                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)
                if cfg.MRCNN.FUSION:
                    mask_pred = mask_feat
                elif cfg.MODEL.BOUNDARY_ON:
                    mask_pred = mask_feat[0]
                    boundary_pred = mask_feat[1]
                else:
                    mask_pred = 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.BOUNDARY_ON:
                    loss_boundary = mask_rcnn_heads.mask_rcnn_losses_balanced(boundary_pred, rpn_ret['boundary_int32'])
                    return_dict['losses']['loss_boundary'] = loss_boundary

            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
        else:
            # Testing
            return_dict['rois'] = rpn_ret['rois']
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_pred

        return return_dict
示例#18
0
    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)
        # logging.info(f"roi belong to which image: shape {rpn_ret['rois'][:, 0:1].shape}\
        #  \n {rpn_ret['rois'][:, 0]}")
        # 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)
            if self.weak_supervise:
                cls_score, det_score, bbox_pred = self.Box_Outs(box_feat)
            else:
                cls_score, bbox_pred = self.Box_Outs(box_feat)
        else:
            # TODO: complete the returns for RPN only situation
            pass

        if self.training and not self.weak_supervise:
            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)

        elif self.training and self.weak_supervise:
            # Weak supervision image-level loss
            # logging.info(f"image-level labels: shape {rpn_ret['image_labels_vec'].shape}\n {rpn_ret['image_labels_vec']}")
            # logging.info(f"cls score: shape {cls_score.shape}\n {cls_score}")
            # logging.info(f"det score: shape {det_score.shape}\n {det_score}")
            #
            return_dict['losses'] = {}
            return_dict['metrics'] = {}
            image_loss_cls, acc_score, reg = fast_rcnn_heads.image_level_loss(
                cls_score, det_score, rpn_ret['rois'],
                rpn_ret['image_labels_vec'], self.bceloss, box_feat)
            return_dict['losses']['image_loss_cls'] = image_loss_cls
            return_dict['losses']['spatial_reg'] = reg

            return_dict['metrics']['accuracy_cls'] = acc_score

            # 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
示例#19
0
    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

        if cfg.LESION.USE_POSITION:
            blob_conv, res5_feat = self.Conv_Body(im_data)
        elif cfg.LESION.SHALLOW_POSITION:
            blob_conv, pos_cls_pred, pos_reg_pred = self.Conv_Body(im_data)
        else:
            blob_conv = 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 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)
            # 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
            if cfg.LESION.SHALLOW_POSITION:
                return_dict['pos_cls_pred'] = pos_cls_pred
                return_dict['pos_reg_pred'] = pos_reg_pred
            if cfg.LESION.POSITION_RCNN:
                pos_cls_rcnn, pos_reg_rcnn = self.Position_RCNN(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

            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
            # RCNN Position
            if cfg.LESION.POSITION_RCNN:
                pos_cls_rcnn_loss, pos_reg_rcnn_loss, accuracy_position_rcnn = position_rcnn_losses(
                    pos_cls_rcnn, pos_reg_rcnn, roidb)
                return_dict['losses']['RcnnPosCls_loss'] = pos_cls_rcnn_loss
                return_dict['losses']['RcnnPosReg_loss'] = pos_reg_rcnn_loss
                return_dict['metrics'][
                    'accuracy_position_rcnn'] = accuracy_position_rcnn
            # Shallow Position Branch
            elif cfg.LESION.SHALLOW_POSITION:
                pos_cls_loss, pos_reg_loss, accuracy_position = position_losses(
                    pos_cls_pred, pos_reg_pred, roidb)
                #pos_reg_loss = position_reg_losses(reg_pred, roidb)
                return_dict['losses']['pos_cls_loss'] = pos_cls_loss
                return_dict['losses']['pos_reg_loss'] = pos_reg_loss
                return_dict['metrics']['accuracy_position'] = accuracy_position
            # Position Branch
            elif 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)
                pos_cls_loss, pos_reg_loss, accuracy_position = position_losses(
                    pos_cls_pred, pos_reg_pred, roidb)
                #pos_reg_loss = position_reg_losses(reg_pred, roidb)
                return_dict['losses']['pos_cls_loss'] = pos_cls_loss
                return_dict['losses']['pos_reg_loss'] = pos_reg_loss
                return_dict['metrics']['accuracy_position'] = accuracy_position

            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
示例#20
0
    def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
        M = cfg.LESION.NUM_IMAGES_3DCE
        # data.shape: [n,M*3,h,w]
        n, c, h, w = data.shape
        im_data = data.view(n * M, 3, h, w)
        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)

        # blob.shape [n, c,h,w] for 2d
        # blob.shape [nM,c,h,w] for 3DCE

        blob_conv_for_RPN = []
        blob_conv_for_RCNN = []
        # 12/25,concat all slices before RPN.
        if cfg.LESION.CONCAT_BEFORE_RPN:
            for blob in blob_conv:
                _, c, h, w = blob.shape
                blob = blob.view(n, M * c, h, w)
                blob_conv_for_RPN.append(blob)
            blob_conv_for_RCNN = blob_conv_for_RPN
        # 01/20,ele-sum all slices before RPN.
        elif cfg.LESION.SUM_BEFORE_RPN:
            for blob in blob_conv:
                blob_ = 0
                for i in range(M):
                    blob_ += blob[n * i:n * (i + 1), :, :, :]
                blob_conv_for_RPN.append(blob_)
                _, c, h, w = blob.shape
                blob = blob.view(n, M * c, h, w)
                blob_conv_for_RCNN.append(blob)
        # Only support three_slices each modality currently.
        elif cfg.LESION.MULTI_MODALITY:
            for blob in blob_conv:
                _, c, h, w = blob.shape
                blob = blob.view(n, M * c, h, w)
                m1_blob_conv = blob[:, 0:c, :, :]
                m2_blob_conv = blob[:, c:, :, :]
                blob_conv_for_RPN.append(
                    self.gif_net(m1_blob_conv, m2_blob_conv))
            blob_conv_for_RCNN = blob_conv_for_RPN
        # Standard 3DCE, feed middle slice into RPN.
        else:
            for blob in blob_conv:
                _, c, h, w = blob.shape
                blob = blob.view(n, M * c, h, w)
                blob_conv_for_RPN.append(blob[:, (M // 2) * c:(M // 2 + 1) *
                                              c, :, :])
                blob_conv_for_RCNN.append(blob)

        rpn_ret = self.RPN(blob_conv_for_RPN, 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_for_RCNN = blob_conv_for_RCNN[-self.num_roi_levels:]

        if not self.training:
            return_dict['blob_conv'] = blob_conv_for_RCNN

        if not cfg.MODEL.RPN_ONLY:
            if cfg.MODEL.SHARE_RES5 and self.training:
                box_feat, res5_feat = self.Box_Head(blob_conv_for_RCNN,
                                                    rpn_ret)
            else:
                box_feat = self.Box_Head(blob_conv_for_RCNN, rpn_ret)
            cls_score, bbox_pred = self.Box_Outs(box_feat)
            # print cls_score.shape
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_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

            # 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
示例#21
0
    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_accuray_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_accuray_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
示例#22
0
    def _forward(self, data, im_info, roidb=None, **rpn_kwargs):
        # data.shape: [n,c,cfg.LESION.NUM_SLICE,h,w]
        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[i].shape [n,c,d,h,w] for RPN, d = cfg.LESION.SLICE_NUM
        blob_conv_for_RPN,blob_conv_for_RCNN = self.Conv_Body(im_data)

        # blob.shape [n,c,h,w] for RPN
        # blob.shape [n,c,h,w] for RCNN
        #blob_conv_for_RPN = []
        #blob_conv_for_RCNN = []
        #if RPN_2_RCNN_3:
        #    for blob in blob_conv:
        #        n, c, d, h, w = blob.shape
        #        blob_ = blob[:, :, d//2, :, :]
        #        blob_conv_for_RPN.append(blob_)
        #    blob_conv_for_RCNN = blob_conv


        rpn_ret = self.RPN(blob_conv_for_RPN, 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_for_RCNN = blob_conv_for_RCNN[-self.num_roi_levels:]

        if not self.training:
            return_dict['blob_conv'] = blob_conv_for_RCNN

        if not cfg.MODEL.RPN_ONLY:
            if cfg.MODEL.SHARE_RES5 and self.training:
                box_feat, res5_feat = self.Box_Head(blob_conv_for_RCNN, rpn_ret)
            else:
                box_feat = self.Box_Head(blob_conv_for_RCNN, rpn_ret)
            cls_score, bbox_pred = self.Box_Outs(box_feat)
            # print cls_score.shape
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_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


            # 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))

        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 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

            # we only use the car cls
            car_cls_int = 4
            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'] == car_cls_int)
                if len(cfg.TRAIN.CE_CAR_CLS_FINETUNE_WIGHT):
                    ce_weight = np.array(cfg.TRAIN.CE_CAR_CLS_FINETUNE_WIGHT)
                else:
                    ce_weight = None

                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=self.shape_sim_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(), self.shape_sim_mat,
                    rpn_ret['car_cls_labels_int32'][car_idx].astype('int64'))
                return_dict['metrics']['rot_diff_degree'] = rot_sim(
                    rot_pred[car_idx].data.cpu().numpy(),
                    rpn_ret['quaternions'][car_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)
                car_idx = np.where(rpn_ret['labels_int32'] == car_cls_int)

                # Build translation head heres from the bounding box
                if cfg.TRANS_HEAD.INPUT_CONV_BODY:
                    pred_boxes_car = pred_boxes[:, 4 * car_cls_int:4 *
                                                (car_cls_int + 1)].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[car_idx]
                elif cfg.TRANS_HEAD.INPUT_TRIPLE_HEAD:
                    pred_boxes_car = pred_boxes[:, 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, car_cls_rot_feat)
                    car_trans_pred = car_trans_pred[car_idx]
                else:
                    pred_boxes_car = pred_boxes[car_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'][car_idx]
                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'][car_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'][car_idx].astype(
                        'int64')
                else:
                    # Using the predicted car id
                    print("Not properly implemented for pytorch")
                    car_ids = car_cls_score[car_idx].max(dim=1)
                # Get mesh vertices and generate loss
                UV_projection_loss = plane_projection_loss(
                    car_trans_pred, label_trans, rot_pred[car_idx],
                    rpn_ret['quaternions'][car_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
示例#24
0
    def _forward(self, data, rois, labels):
        im_data = data
        if self.training:
            rois = rois.squeeze(dim=0).type(im_data.dtype)
            labels = labels.squeeze(dim=0).type(im_data.dtype)

        return_dict = {}  # A dict to collect return variables

        blob_conv = self.Conv_Body(im_data).contiguous()

        if not self.training:
            return_dict['blob_conv'] = blob_conv

        box_feat = self.Box_Head(blob_conv, rois)
        mil_score = self.Box_MIL_Outs(box_feat)
        refine_score = self.Box_Refine_Outs(box_feat)
        if cfg.MODEL.WITH_FRCNN:
            cls_score, bbox_pred = self.FRCNN_Outs(box_feat)

        device = box_feat.device

        if self.training:
            return_dict['losses'] = {}

            # image classification loss
            im_cls_score = mil_score.sum(dim=0, keepdim=True)
            loss_im_cls = pcl_heads.mil_losses(im_cls_score, labels)
            return_dict['losses']['loss_im_cls'] = loss_im_cls

            # refinement loss
            boxes = rois.data.cpu().numpy()
            im_labels = labels.data.cpu().numpy()
            boxes = boxes[:, 1:]

            for i_refine, refine in enumerate(refine_score):
                if i_refine == 0:
                    pcl_output = OICR(boxes, mil_score, im_labels, refine)
                else:
                    pcl_output = OICR(boxes, refine_score[i_refine - 1],
                                      im_labels, refine)

                refine_loss = self.Refine_Losses[i_refine](
                    refine, Variable(torch.from_numpy(
                        pcl_output['labels'])).to(device),
                    Variable(torch.from_numpy(
                        pcl_output['cls_loss_weights'])).to(device),
                    Variable(torch.from_numpy(
                        pcl_output['gt_assignment'])).to(device))

                if i_refine == 0:
                    refine_loss *= 3.0

                return_dict['losses']['refine_loss%d' %
                                      i_refine] = refine_loss.clone()

            if cfg.MODEL.WITH_FRCNN:
                labels, cls_loss_weights, bbox_targets, bbox_inside_weights, \
                    bbox_outside_weights = fast_rcnn_heads.get_fast_rcnn_targets(
                        boxes, refine_score, im_labels)

                cls_loss, bbox_loss = fast_rcnn_heads.fast_rcnn_losses(
                    cls_score, bbox_pred,
                    Variable(torch.from_numpy(labels)).to(device),
                    Variable(torch.from_numpy(cls_loss_weights)).to(device),
                    Variable(torch.from_numpy(bbox_targets)).to(device),
                    Variable(torch.from_numpy(bbox_inside_weights)).to(device),
                    Variable(
                        torch.from_numpy(bbox_outside_weights)).to(device))

                return_dict['losses']['cls_loss'] = cls_loss
                return_dict['losses']['bbox_loss'] = bbox_loss

            # pytorch0.4 bug on gathering scalar(0-dim) tensors
            for k, v in return_dict['losses'].items():
                return_dict['losses'][k] = v.unsqueeze(0)

        else:
            # Testing
            return_dict['rois'] = rois
            return_dict['mil_score'] = mil_score
            return_dict['refine_score'] = refine_score
            if cfg.MODEL.WITH_FRCNN:
                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)

        rois = Variable(torch.from_numpy(rpn_ret['rois'])).cuda(device_id)
        return_dict['rois'] = rois
        if self.training:
            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)
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_pred
        else:
            # TODO: complete the returns for RPN only situation
            pass

        if self.training:
            # 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)
            return_dict['loss_rpn_cls'] = loss_rpn_cls
            return_dict['loss_rpn_bbox'] = loss_rpn_bbox

            # bbox loss
            loss_cls, loss_bbox = 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['loss_rcnn_cls'] = loss_cls
            return_dict['loss_rcnn_bbox'] = loss_bbox

            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['loss_rcnn_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['loss_rcnn_keypoints'] = loss_keypoints

        return return_dict
示例#26
0
    def _forward(self, data, query, im_info, query_type, roidb=None, **rpn_kwargs):

        #query_type = query_type.item()

        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

        # feed image data to base model to obtain base feature map
        blob_conv = self.Conv_Body(im_data)

        query_conv = []
        shot = len(query)
        for i in range(shot):
            query_conv.append(self.Conv_Body(query[i]))
        
        def pooling(feats, method='avg', dim = 0):
            feats = torch.stack(feats)
            if method == 'avg':
                feat = torch.mean(feats, dim=dim)
            elif method == 'max':
                feat, _ = torch.max(feats, dim=dim)
            return feat
        
        if cfg.FPN.FPN_ON:
            query_conv = list(map(list, zip(*query_conv)))
            query_conv = [pooling(QC_i, method='avg', dim = 0) for QC_i in query_conv]

            rpn_feat = []
            act_feat = []
            act_aim = []
            c_weight = []
            
            if len(blob_conv) == 5:
                start_match_idx = 1
                #sa_blob_conv = self.sa(blob_conv[0])
                sa_blob_conv = blob_conv[0]
                rpn_feat.append(sa_blob_conv)
            else:
                start_match_idx = 0

            for IP, QP in zip(blob_conv[start_match_idx:], query_conv[start_match_idx:]):
                _rpn_feat, _act_feat, _act_aim, _c_weight = self.match_net(IP, QP)
                rpn_feat.append(_rpn_feat)
                act_feat.append(_act_feat)
                act_aim.append(_act_aim)
                c_weight.append(_c_weight)
                """
                correlation = []
                QP_pool = self.global_avgpool(QP)
                #IP_sa = self.sa(IP)
                IP_sa = IP
                for IP_sa_batch, QP_pool_batch in zip(IP_sa, QP_pool):
                    IP_sa_batch, QP_pool_batch = IP_sa_batch.unsqueeze(0), QP_pool_batch.unsqueeze(0)
                    correlation.append(F.conv2d(IP_sa_batch, QP_pool_batch.permute(1,0,2,3), groups=IP.shape[1]).squeeze(0))
                correlation = torch.stack(correlation)
                rpn_feat.append(correlation)
                act_feat.append(correlation)
                act_aim.append(QP)
                """
        else:
            query_conv = pooling(query_conv)
            rpn_feat, act_feat, act_aim, c_weight = self.match_net(blob_conv, query_conv)
            """
            correlation = []
            QP_pool = self.global_avgpool(QP)
            #IP_sa = self.sa(IP)
            IP_sa = IP
            for IP_sa_batch, QP_pool_batch in zip(IP_sa, QP_pool):
                IP_sa_batch, QP_pool_batch = IP_sa_batch.unsqueeze(0), QP_pool_batch.unsqueeze(0)
                correlation.append(F.conv2d(IP_sa_batch, QP_pool_batch.permute(1,0,2,3), groups=IP.shape[1]).squeeze(0))
            correlation = torch.stack(correlation)
            rpn_feat = correlation
            act_feat = correlation
            act_aim = query_conv
            """
        rpn_ret = self.RPN(rpn_feat, 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
            return_dict['query_conv'] = query_conv

        if not cfg.MODEL.RPN_ONLY:
            if not cfg.FPN.FPN_ON:
                if cfg.MODEL.SHARE_RES5 and self.training:
                    if cfg.RELATION_RCNN:
                        box_feat, query_box_feat, res5_feat, query_res5_feat = self.Box_Head(act_feat, act_aim, rpn_ret)
                    else:
                        box_feat, res5_feat = self.Box_Head(act_feat, rpn_ret)
                else:
                    if cfg.RELATION_RCNN:
                        box_feat, query_box_feat= self.Box_Head(act_feat, act_aim, rpn_ret)
                    else:
                        box_feat = self.Box_Head(act_feat, rpn_ret)
                
                if cfg.RELATION_RCNN:
                    cls_score, bbox_pred = self.Box_Outs(box_feat, query_box_feat)
                else:
                    cls_score, bbox_pred = self.Box_Outs(box_feat)
            else:
                if cfg.RELATION_RCNN:
                    box_feat, query_box_feat = self.Box_Head(act_feat, act_aim, rpn_ret)
                    cls_score, bbox_pred = self.Box_Outs(box_feat, query_box_feat)
                else:
                    box_feat = self.Box_Head(act_feat, 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'))
            ))
            #rpn_kwargs.update({'query_type': query_type})
            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'], rpn_ret['rois'], query_type, use_marginloss=False)
            #return_dict['losses']['margin_loss'] = margin_loss
            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):
                    if cfg.RELATION_RCNN:
                        mask_feat = self.Mask_Head(res5_feat, query_res5_feat, rpn_ret,
                                                roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
                    else:
                        mask_feat = self.Mask_Head(res5_feat, rpn_ret, roi_has_mask_int32=rpn_ret['roi_has_mask_int32'])
                else:
                    if cfg.RELATION_RCNN:
                        mask_feat = self.Mask_Head(act_feat, act_aim, rpn_ret)
                    else:
                        mask_feat = self.Mask_Head(act_feat, 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'], rpn_ret['mask_rois'], query_type)
                #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)
            
            return_dict['rois'] = rpn_ret['rois']
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_pred

        else:
            # Testing
            return_dict['rois'] = rpn_ret['rois']
            return_dict['cls_score'] = cls_score
            return_dict['bbox_pred'] = bbox_pred

        return return_dict
示例#27
0
    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