aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
Diffstat
-rw-r--r--tensorflow/core/framework/shape_inference.cc3
-rw-r--r--tensorflow/core/kernels/non_max_suppression_op.cc2
-rw-r--r--tensorflow/core/ops/image_ops.cc77
-rw-r--r--tensorflow/python/ops/image_ops_test.py35
4 files changed, 71 insertions, 46 deletions
diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc
index 8d597e198d..3e77028a5f 100644
--- a/tensorflow/core/framework/shape_inference.cc
+++ b/tensorflow/core/framework/shape_inference.cc
@@ -950,8 +950,7 @@ Status InferenceContext::GetScalarFromTensor(const Tensor* t, int64* val) {
*val = t->scalar<int64>()();
return Status::OK();
} else {
- return errors::InvalidArgument(
- "Scalar input for dim size must be int32 or int64");
+ return errors::InvalidArgument("Scalar input must be int32 or int64.");
}
}
diff --git a/tensorflow/core/kernels/non_max_suppression_op.cc b/tensorflow/core/kernels/non_max_suppression_op.cc
index c7d0d4de0d..5d9257e20b 100644
--- a/tensorflow/core/kernels/non_max_suppression_op.cc
+++ b/tensorflow/core/kernels/non_max_suppression_op.cc
@@ -126,7 +126,7 @@ void DoNonMaxSuppressionOp(
const Tensor& max_output_size, const float score_threshold,
const std::function<bool(int, int)>& suppress_check_fn,
bool pad_to_max_output_size = false, int* ptr_num_valid_outputs = nullptr) {
- const int output_size = std::min(max_output_size.scalar<int>()(), num_boxes);
+ const int output_size = max_output_size.scalar<int>()();
std::vector<float> scores_data(num_boxes);
std::copy_n(scores.flat<float>().data(), num_boxes, scores_data.begin());
diff --git a/tensorflow/core/ops/image_ops.cc b/tensorflow/core/ops/image_ops.cc
index 81f324a3ef..11ca0bd259 100644
--- a/tensorflow/core/ops/image_ops.cc
+++ b/tensorflow/core/ops/image_ops.cc
@@ -108,6 +108,29 @@ Status ColorspaceShapeFn(InferenceContext* c) {
return Status::OK();
}
+Status NMSShapeFn(InferenceContext* c) {
+ // Get inputs and validate ranks.
+ ShapeHandle boxes;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &boxes));
+ ShapeHandle scores;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &scores));
+ ShapeHandle max_output_size;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &max_output_size));
+ ShapeHandle iou_threshold;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &iou_threshold));
+ ShapeHandle score_threshold;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &score_threshold));
+ // The boxes is a 2-D float Tensor of shape [num_boxes, 4].
+ DimensionHandle unused;
+ // The boxes[0] and scores[0] are both num_boxes.
+ TF_RETURN_IF_ERROR(c->Merge(c->Dim(boxes, 0), c->Dim(scores, 0), &unused));
+ // The boxes[1] is 4.
+ TF_RETURN_IF_ERROR(c->WithValue(c->Dim(boxes, 1), 4, &unused));
+
+ c->set_output(0, c->Vector(c->UnknownDim()));
+ return Status::OK();
+}
+
} // namespace
// --------------------------------------------------------------------------
@@ -694,29 +717,7 @@ REGISTER_OP("NonMaxSuppressionV3")
.Input("iou_threshold: float")
.Input("score_threshold: float")
.Output("selected_indices: int32")
- .SetShapeFn([](InferenceContext* c) {
- // Get inputs and validate ranks.
- ShapeHandle boxes;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &boxes));
- ShapeHandle scores;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &scores));
- ShapeHandle max_output_size;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &max_output_size));
- ShapeHandle iou_threshold;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &iou_threshold));
- ShapeHandle score_threshold;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &score_threshold));
- // The boxes is a 2-D float Tensor of shape [num_boxes, 4].
- DimensionHandle unused;
- // The boxes[0] and scores[0] are both num_boxes.
- TF_RETURN_IF_ERROR(
- c->Merge(c->Dim(boxes, 0), c->Dim(scores, 0), &unused));
- // The boxes[1] is 4.
- TF_RETURN_IF_ERROR(c->WithValue(c->Dim(boxes, 1), 4, &unused));
-
- c->set_output(0, c->Vector(c->UnknownDim()));
- return Status::OK();
- });
+ .SetShapeFn(NMSShapeFn);
REGISTER_OP("NonMaxSuppressionV4")
.Input("boxes: float")
@@ -728,26 +729,16 @@ REGISTER_OP("NonMaxSuppressionV4")
.Output("valid_outputs: int32")
.Attr("pad_to_max_output_size: bool = false")
.SetShapeFn([](InferenceContext* c) {
- // Get inputs and validate ranks.
- ShapeHandle boxes;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &boxes));
- ShapeHandle scores;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &scores));
- ShapeHandle max_output_size;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &max_output_size));
- ShapeHandle iou_threshold;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &iou_threshold));
- ShapeHandle score_threshold;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &score_threshold));
- // The boxes is a 2-D float Tensor of shape [num_boxes, 4].
- DimensionHandle unused;
- // The boxes[0] and scores[0] are both num_boxes.
- TF_RETURN_IF_ERROR(
- c->Merge(c->Dim(boxes, 0), c->Dim(scores, 0), &unused));
- // The boxes[1] is 4.
- TF_RETURN_IF_ERROR(c->WithValue(c->Dim(boxes, 1), 4, &unused));
-
- c->set_output(0, c->Vector(c->UnknownDim()));
+ TF_RETURN_IF_ERROR(NMSShapeFn(c));
+
+ bool pad_to_max;
+ TF_RETURN_IF_ERROR(c->GetAttr("pad_to_max_output_size", &pad_to_max));
+ if (pad_to_max) {
+ // If padded, overwrite the shape of the output to be static.
+ DimensionHandle output_dim;
+ TF_RETURN_IF_ERROR(c->MakeDimForScalarInput(2, &output_dim));
+ c->set_output(0, c->MakeShape({output_dim}));
+ }
c->set_output(1, c->MakeShape({}));
return Status::OK();
});
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 0e4193e23b..2c61bb232a 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -3658,6 +3658,41 @@ class NonMaxSuppressionTest(test_util.TensorFlowTestCase):
image_ops.non_max_suppression(boxes, scores, 3, [[0.5]])
+class NonMaxSuppressionPaddedTest(test_util.TensorFlowTestCase):
+
+ def testSelectFromThreeClusters(self):
+ boxes_np = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
+ [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
+ scores_np = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
+ max_output_size_np = 5
+ iou_threshold_np = 0.5
+ boxes = constant_op.constant(boxes_np)
+ scores = constant_op.constant(scores_np)
+ max_output_size = constant_op.constant(max_output_size_np)
+ iou_threshold = constant_op.constant(iou_threshold_np)
+ selected_indices_padded, num_valid_padded = \
+ image_ops.non_max_suppression_padded(
+ boxes,
+ scores,
+ max_output_size,
+ iou_threshold,
+ pad_to_max_output_size=True)
+ selected_indices, num_valid = image_ops.non_max_suppression_padded(
+ boxes,
+ scores,
+ max_output_size,
+ iou_threshold,
+ pad_to_max_output_size=False)
+ # The output shape of the padded operation must be fully defined.
+ self.assertEqual(selected_indices_padded.shape.is_fully_defined(), True)
+ self.assertEqual(selected_indices.shape.is_fully_defined(), False)
+ with self.test_session():
+ self.assertAllClose(selected_indices_padded.eval(), [3, 0, 5, 0, 0])
+ self.assertEqual(num_valid_padded.eval(), 3)
+ self.assertAllClose(selected_indices.eval(), [3, 0, 5])
+ self.assertEqual(num_valid.eval(), 3)
+
+
class VerifyCompatibleImageShapesTest(test_util.TensorFlowTestCase):
"""Tests utility function used by ssim() and psnr()."""
generated by cgit v1.2.3 (git 2.39.1) at 2024-06-27 01:33:38 +0000