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)