def ArgSortJob(
input: oft.ListNumpy.Placeholder(
tuple([dim + 10 for dim in in_shape]),
dtype=type_name_to_flow_type[data_type],
)
):
with flow.scope.placement(device_type, "0:0"):
return flow.argsort(input, direction)
def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
scores_inds = flow_exp.argsort(scores, dim=0, descending=True)
boxes = flow._C.gather(boxes, scores_inds, axis=0)
_nms_op = (flow_exp.builtin_op("nms").Input("in").Output("out").Attr(
"iou_threshold", iou_threshold).Attr("keep_n", -1).Build())
keep = _nms_op(boxes)[0]
index = flow_exp.squeeze(flow_exp.argwhere(keep), dim=[1])
return flow._C.gather(scores_inds, index, axis=0)
def _test_argsort(test_case, data_shape, axis, descending, data_type, device):
input = flow.tensor(
np.random.randn(*data_shape),
dtype=type_name_to_flow_type[data_type],
device=flow.device(device),
)
of_out = flow.argsort(input, dim=axis, descending=descending)
np_input = -input.numpy() if descending else input.numpy()
np_out = np.argsort(np_input, axis=axis)
test_case.assertTrue(np.array_equal(of_out.numpy().flatten(), np_out.flatten()))
def _argsort(self, dim=None, descending=None):
return flow.argsort(self, dim=dim, descending=descending)
def nms_op(boxes, scores, iou_threshold: float):
score_inds = flow.argsort(scores, dim=0, descending=True)
boxes = flow._C.gather(boxes, score_inds, axis=0)
keep = flow._C.nms(boxes, iou_threshold)
index = flow.squeeze(flow.argwhere(keep), dim=[1])
return flow._C.gather(score_inds, index, axis=0)
