27 #include "tensorflow/core/framework/op.h" 28 #include "tensorflow/core/framework/op_kernel.h" 29 #include "tensorflow/core/lib/core/errors.h" 35 class OutputAllocator {
37 OutputAllocator(tensorflow::OpKernelContext* context) : context(context) {}
39 void AllocKeepIndices(int64_t** ptr, int64_t num) {
43 TensorShape shape({num});
44 OP_REQUIRES_OK(context, context->allocate_output(0, shape, &tensor));
45 auto flat_tensor = tensor->flat<int64>();
46 *ptr = (int64_t*)flat_tensor.data();
50 tensorflow::OpKernelContext* context;
54 class NmsOpKernel :
public tensorflow::OpKernel {
56 explicit NmsOpKernel(tensorflow::OpKernelConstruction* construction)
57 : OpKernel(construction) {
58 OP_REQUIRES_OK(construction,
59 construction->GetAttr(
"nms_overlap_thresh",
60 &nms_overlap_thresh));
63 void Compute(tensorflow::OpKernelContext* context)
override {
65 const Tensor& boxes = context->input(0);
66 const Tensor& scores = context->input(1);
70 Dim num_points(
"num_points");
76 Kernel(context, boxes, scores);
80 virtual void Kernel(tensorflow::OpKernelContext* context,
81 const tensorflow::Tensor& boxes,
82 const tensorflow::Tensor& scores) = 0;
85 float nms_overlap_thresh;
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:204
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:69
Definition: ShapeChecking.h:35