def forward(self, imgs, bboxs, scale):
n,_,_ = bboxs.size()
if n != 1:
raise ValueError('Currently only batch size 1 is supported.')
_, _, H, W = imgs.size()
img_size = (H, W)
features = self.head_detector.extractor(imgs)
rpn_locs, rpn_scores, rois, rois_scores, anchor = self.head_detector.rpn(features, img_size, scale)
bbox = bboxs[0]
rpn_score = rpn_scores[0]
rpn_loc = rpn_locs[0]
# ------------------ RPN losses -------------------#
gt_rpn_loc, gt_rpn_label = self.anchor_target_creator(at.tonumpy(bbox),anchor,img_size)
gt_rpn_label = at.tovariable(gt_rpn_label).long()
gt_rpn_loc = at.tovariable(gt_rpn_loc)
rpn_loc_loss = head_detector_loss(
rpn_loc,
gt_rpn_loc,
gt_rpn_label.data,
self.rpn_sigma)
rpn_cls_loss = F.cross_entropy(rpn_score, gt_rpn_label.cuda(), ignore_index=-1)
_gt_rpn_label = gt_rpn_label[gt_rpn_label > -1]
_rpn_score = at.tonumpy(rpn_score)[at.tonumpy(gt_rpn_label) > -1]
self.rpn_cm.add(at.totensor(_rpn_score, False), _gt_rpn_label.data.long())
losses = [rpn_loc_loss, rpn_cls_loss]
losses = losses + [sum(losses)]
return LossTuple(*losses), rois, rois_scores
def predict(self, x, scale=1., mode='evaluate', thresh=0.01):
if mode == 'evaluate':
nms_thresh = 0.3 #0.3
score_thresh = thresh
elif mode == 'visualize':
nms_thresh = 0.3 #0.3
score_thresh = thresh
_, _, rois, rois_scores, _ = self.forward(x, scale=scale)
roi = at.totensor(rois)
probabilities = at.tonumpy(F.softmax(at.tovariable(rois_scores)))
_, _, H, W = x.size()
size = (H,W)
roi[:, 0::2] = (roi[:, 0::2]).clamp(min=0, max=size[0])
roi[:, 1::2] = (roi[:, 1::2]).clamp(min=0, max=size[1])
roi_raw = at.tonumpy(roi)
probabilities = np.squeeze(probabilities)
bbox, score = self._suppress(roi_raw, probabilities, nms_thresh, score_thresh)
return bbox, score