| 132 | } |
| 133 | |
| 134 | Status SoftNMSShapeFn(InferenceContext* c) { |
| 135 | // Get inputs and validate ranks. |
| 136 | ShapeHandle boxes; |
| 137 | TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &boxes)); |
| 138 | ShapeHandle scores; |
| 139 | TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &scores)); |
| 140 | ShapeHandle max_output_size; |
| 141 | TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &max_output_size)); |
| 142 | ShapeHandle iou_threshold; |
| 143 | TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &iou_threshold)); |
| 144 | ShapeHandle score_threshold; |
| 145 | TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &score_threshold)); |
| 146 | ShapeHandle soft_nms_sigma; |
| 147 | TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &soft_nms_sigma)); |
| 148 | // The boxes is a 2-D float Tensor of shape [num_boxes, 4]. |
| 149 | DimensionHandle unused; |
| 150 | // The boxes[0] and scores[0] are both num_boxes. |
| 151 | TF_RETURN_IF_ERROR(c->Merge(c->Dim(boxes, 0), c->Dim(scores, 0), &unused)); |
| 152 | // The boxes[1] is 4. |
| 153 | TF_RETURN_IF_ERROR(c->WithValue(c->Dim(boxes, 1), 4, &unused)); |
| 154 | |
| 155 | c->set_output(0, c->Vector(c->UnknownDim())); |
| 156 | c->set_output(1, c->Vector(c->UnknownDim())); |
| 157 | return Status::OK(); |
| 158 | } |
| 159 | |
| 160 | Status CombinedNMSShapeFn(InferenceContext* c) { |
| 161 | // Get inputs and validate ranks |
no test coverage detected