00001
00002
00003
00004
00005
00006
00007
00008 """
00009 Proposal Target Operator selects foreground and background roi and assigns label, bbox_transform to them.
00010 """
00011
00012 import mxnet as mx
00013 import numpy as np
00014 from distutils.util import strtobool
00015
00016
00017
00018
00019 class BoxAnnotatorOHEMOperator(mx.operator.CustomOp):
00020 def __init__(self, num_classes, num_reg_classes, roi_per_img):
00021 super(BoxAnnotatorOHEMOperator, self).__init__()
00022 self._num_classes = num_classes
00023 self._num_reg_classes = num_reg_classes
00024 self._roi_per_img = roi_per_img
00025
00026 def forward(self, is_train, req, in_data, out_data, aux):
00027
00028 cls_score = in_data[0]
00029 bbox_pred = in_data[1]
00030 labels = in_data[2].asnumpy()
00031 bbox_targets = in_data[3]
00032 bbox_weights = in_data[4]
00033
00034 per_roi_loss_cls = mx.nd.SoftmaxActivation(cls_score) + 1e-14
00035 per_roi_loss_cls = per_roi_loss_cls.asnumpy()
00036 per_roi_loss_cls = per_roi_loss_cls[np.arange(per_roi_loss_cls.shape[0], dtype='int'), labels.astype('int')]
00037 per_roi_loss_cls = -1 * np.log(per_roi_loss_cls)
00038 per_roi_loss_cls = np.reshape(per_roi_loss_cls, newshape=(-1,))
00039
00040 per_roi_loss_bbox = bbox_weights * mx.nd.smooth_l1((bbox_pred - bbox_targets), scalar=1.0)
00041 per_roi_loss_bbox = mx.nd.sum(per_roi_loss_bbox, axis=1).asnumpy()
00042
00043 top_k_per_roi_loss = np.argsort(per_roi_loss_cls + per_roi_loss_bbox)
00044 labels_ohem = labels
00045 labels_ohem[top_k_per_roi_loss[::-1][self._roi_per_img:]] = -1
00046 bbox_weights_ohem = bbox_weights.asnumpy()
00047 bbox_weights_ohem[top_k_per_roi_loss[::-1][self._roi_per_img:]] = 0
00048
00049 labels_ohem = mx.nd.array(labels_ohem)
00050 bbox_weights_ohem = mx.nd.array(bbox_weights_ohem)
00051
00052 for ind, val in enumerate([labels_ohem, bbox_weights_ohem]):
00053 self.assign(out_data[ind], req[ind], val)
00054
00055
00056 def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
00057 for i in range(len(in_grad)):
00058 self.assign(in_grad[i], req[i], 0)
00059
00060
00061 @mx.operator.register('BoxAnnotatorOHEM')
00062 class BoxAnnotatorOHEMProp(mx.operator.CustomOpProp):
00063 def __init__(self, num_classes, num_reg_classes, roi_per_img):
00064 super(BoxAnnotatorOHEMProp, self).__init__(need_top_grad=False)
00065 self._num_classes = int(num_classes)
00066 self._num_reg_classes = int(num_reg_classes)
00067 self._roi_per_img = int(roi_per_img)
00068
00069 def list_arguments(self):
00070 return ['cls_score', 'bbox_pred', 'labels', 'bbox_targets', 'bbox_weights']
00071
00072 def list_outputs(self):
00073 return ['labels_ohem', 'bbox_weights_ohem']
00074
00075 def infer_shape(self, in_shape):
00076 labels_shape = in_shape[2]
00077 bbox_weights_shape = in_shape[4]
00078
00079 return in_shape, \
00080 [labels_shape, bbox_weights_shape]
00081
00082 def create_operator(self, ctx, shapes, dtypes):
00083 return BoxAnnotatorOHEMOperator(self._num_classes, self._num_reg_classes, self._roi_per_img)
00084
00085 def declare_backward_dependency(self, out_grad, in_data, out_data):
00086 return []