1 /**
2 * Copyright 2023 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "nnacl/kernel/non_max_suppression.h"
18 #include "nnacl/kernel/default_kernel_base.h"
19 #include "nnacl/non_max_suppression_parameter.h"
20 #include "nnacl/fp32/non_max_suppression_fp32.h"
21
NonMaxSuppressioExpandDims(int * dst_shape,int * origin_shape,size_t size)22 void NonMaxSuppressioExpandDims(int *dst_shape, int *origin_shape, size_t size) {
23 int i = 0;
24 for (; i < size; i++) {
25 dst_shape[i] = 1;
26 }
27 for (; i < Num3; i++) {
28 dst_shape[i] = origin_shape[i - size];
29 }
30 }
31
NonMaxSuppressionGetParams(NonMaxSuppressionStruct * nm_suppression)32 void NonMaxSuppressionGetParams(NonMaxSuppressionStruct *nm_suppression) {
33 // optional input order: max_output_per_class, iou_threshold, score_threshold
34 nm_suppression->max_output_per_class_ = 0;
35 if (nm_suppression->base_.in_size_ >= Num3) {
36 TensorC *max_output_tensor = nm_suppression->base_.in_[Index3];
37 if (max_output_tensor != NULL && max_output_tensor->data_ != NULL) {
38 nm_suppression->max_output_per_class_ = *(int *)(max_output_tensor->data_);
39 }
40 }
41
42 nm_suppression->iou_threshold_ = 0.0f;
43 if (nm_suppression->base_.in_size_ >= Num4) {
44 TensorC *iou_threshold_tensor = nm_suppression->base_.in_[Index4];
45 if (iou_threshold_tensor != NULL && iou_threshold_tensor->data_ != NULL) {
46 nm_suppression->iou_threshold_ = *(float *)(iou_threshold_tensor->data_);
47 }
48 }
49
50 nm_suppression->score_threshold_ = 0.0f;
51 if (nm_suppression->base_.in_size_ >= Num5) {
52 TensorC *score_threshold_tensor = nm_suppression->base_.in_[Index5];
53 if (score_threshold_tensor != NULL && score_threshold_tensor->data_ != NULL) {
54 nm_suppression->score_threshold_ = *(float *)(score_threshold_tensor->data_);
55 }
56 }
57 }
58
NonMaxSuppressionCompute(KernelBase * self)59 int NonMaxSuppressionCompute(KernelBase *self) {
60 NonMaxSuppressionStruct *nm_suppression = (NonMaxSuppressionStruct *)self;
61 NNACL_CHECK_NULL_RETURN_ERR(nm_suppression);
62
63 NonMaxSuppressionGetParams(nm_suppression);
64
65 TensorC *box_tensor = self->in_[Index0];
66 NNACL_CHECK_NULL_RETURN_ERR(box_tensor);
67 int box_dims[Num3] = {0}; // batch, box_num, 4
68 bool simple_out = false;
69 if (box_tensor->shape_size_ != Num3) {
70 NonMaxSuppressioExpandDims(box_dims, box_tensor->shape_, Num3 - box_tensor->shape_size_);
71 simple_out = true;
72 }
73 if (box_dims[Index2] != Num4) {
74 return NNACL_NON_MAX_SUPPRESSION_BOX_DIMS_INVALID;
75 }
76
77 TensorC *score_tensor = self->in_[Index1];
78 NNACL_CHECK_NULL_RETURN_ERR(score_tensor);
79 int score_dims[Num3] = {0}; // batch, class, box_num
80 if (score_tensor->shape_size_ != Num3) {
81 NonMaxSuppressioExpandDims(score_dims, score_tensor->shape_, Num3 - score_tensor->shape_size_);
82 }
83 if (score_dims[Index0] != box_dims[Index0]) {
84 return NNACL_NON_MAX_SUPPRESSION_BOX_DIMS_SCORE_UNMATCH;
85 }
86 if (score_dims[Index2] != box_dims[Index1]) {
87 return NNACL_NON_MAX_SUPPRESSION_DIMENSION_SPATIAL_UNMATCH;
88 }
89 if (nm_suppression->base_.out_[OUTPUT_INDEX]->data_ != NULL) {
90 /* output shape and data set in compute */
91 return NNACL_NON_MAX_SUPPRESSION_UNSUPPORT_DEFINE_DATA;
92 }
93 return NonMaxSuppressionSelecte(nm_suppression, simple_out, score_dims);
94 }
95
NonMaxSuppressionPrepare(KernelBase * self)96 int NonMaxSuppressionPrepare(KernelBase *self) {
97 NonMaxSuppressionStruct *nm_suppression = (NonMaxSuppressionStruct *)self;
98 NNACL_CHECK_NULL_RETURN_ERR(nm_suppression);
99
100 // boxes, scores, max_output_boxes, iou_threshold, score_threshold
101 if (self->in_size_ < Num2 || self->in_size_ > Num5 || self->out_size_ != Num1) {
102 return NNACL_NON_MAX_SUPPRESSION_TENSOR_SIZE_INVALID;
103 }
104
105 NNACL_CHECK_NULL_RETURN_ERR(self->in_[FIRST_INPUT]);
106 NNACL_CHECK_NULL_RETURN_ERR(self->out_[OUTPUT_INDEX]);
107 NMSParameter *nmparam = (NMSParameter *)self->param_;
108 NNACL_CHECK_NULL_RETURN_ERR(nmparam);
109 if (nmparam->center_point_box_ != 0 && nmparam->center_point_box_ != 1) {
110 return NNACL_NON_MAX_SUPPRESSION_PARAM_INVALID;
111 }
112 nm_suppression->center_point_box_ = nmparam->center_point_box_;
113 return NNACL_OK;
114 }
115
CreateNonMaxSuppression(OpParameter * param,int data_type)116 KernelBase *CreateNonMaxSuppression(OpParameter *param, int data_type) {
117 NonMaxSuppressionStruct *non_max_suppression = (NonMaxSuppressionStruct *)malloc(sizeof(NonMaxSuppressionStruct));
118 NNACL_CHECK_NULL_RETURN_NULL(non_max_suppression);
119 non_max_suppression->base_.Release = DefaultRelease;
120 non_max_suppression->base_.Resize = DefaultResize;
121 non_max_suppression->base_.Prepare = NonMaxSuppressionPrepare;
122 non_max_suppression->base_.Compute = NonMaxSuppressionCompute;
123 return (KernelBase *)non_max_suppression;
124 }
125
126 REG_KERNEL_CREATOR(PrimType_NonMaxSuppression, kNumberTypeFloat32, CreateNonMaxSuppression)
127