def _pool_and_predict(
self, indices_and_rois, h_cls_seg, h_locs, gt_roi_labels=None):
# PSROI Pooling
# shape: (n_rois, n_class*2, roi_size, roi_size)
pool_cls_seg = _psroi_pooling_2d_yx(
h_cls_seg, indices_and_rois, self.roi_size, self.roi_size,
self.spatial_scale, group_size=self.group_size,
output_dim=self.n_class*2)
# shape: (n_rois, n_class, 2, roi_size, roi_size)
pool_cls_seg = pool_cls_seg.reshape(
(-1, self.n_class, 2, self.roi_size, self.roi_size))
# shape: (n_rois, 2*4, roi_size, roi_size)
pool_locs = _psroi_pooling_2d_yx(
h_locs, indices_and_rois, self.roi_size, self.roi_size,
self.spatial_scale, group_size=self.group_size,
output_dim=2*4)
# Classfication
# Group Max
# shape: (n_rois, n_class, roi_size, roi_size)
h_cls = pool_cls_seg.transpose((0, 1, 3, 4, 2))
h_cls = F.max(h_cls, axis=4)
# Global pooling (vote)
# shape: (n_rois, n_class)
roi_cls_scores = _global_average_pooling_2d(h_cls)
# Bbox Regression
# shape: (n_rois, 2*4)
roi_cls_locs = _global_average_pooling_2d(pool_locs)
n_rois = roi_cls_locs.shape[0]
roi_cls_locs = roi_cls_locs.reshape((n_rois, 2, 4))
# Mask Regression
# shape: (n_rois, n_class, 2, roi_size, roi_size)
# Group Pick by Score
if gt_roi_labels is None:
max_cls_idx = roi_cls_scores.array.argmax(axis=1)
else:
max_cls_idx = gt_roi_labels
# shape: (n_rois, 2, roi_size, roi_size)
roi_seg_scores = pool_cls_seg[np.arange(len(max_cls_idx)), max_cls_idx]
return roi_seg_scores, roi_cls_locs, roi_cls_scores
评论列表
文章目录
