rfcn.py 文件源码

def forward(self, im_data, im_info, gt_boxes=None, gt_ishard=None, dontcare_areas=None):
        features, rois = self.rpn(im_data, im_info, gt_boxes, gt_ishard, dontcare_areas)

        if self.training:
            roi_data = self.proposal_target_layer(rois, gt_boxes, gt_ishard, dontcare_areas, self.n_classes)
            rois = roi_data[0]

        # roi pool
    conv_new1 = self.new_conv(features)
    r_score_map = self.rfcn_score(conv_new1)
    r_bbox_map = self.rfcn_bbox(conv_new1)
    psroi_pooled_cls = self.psroi_pool_cls(r_score_map, rois)
    psroi_pooled_loc = self.psroi_pool_loc(r_bbox_map, rois)
    bbox_pred = self.bbox_pred(psroi_pooled_loc)
    bbox_pred = torch.squeeze(bbox_pred)
    cls_score = self.cls_score(psroi_pooled_cls)
    cls_score = torch.squeeze(cls_score)
        cls_prob = F.softmax(cls_score)

        if self.training:
            self.cross_entropy, self.loss_box = self.build_loss(cls_score, bbox_pred, roi_data)

        return cls_prob, bbox_pred, rois