Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent crash/heap OOB due to integer conversion to unsigned in NMS k…
…ernels

PiperOrigin-RevId: 387938262
Change-Id: Id361a715307e7179977cf5c64391c199a966f2ad
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 31, 2021
1 parent a87fa31 commit 3a73627
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions tensorflow/core/kernels/image/non_max_suppression_op.cc
Expand Up @@ -169,6 +169,8 @@ void DoNonMaxSuppressionOp(OpKernelContext* context, const Tensor& scores,
bool pad_to_max_output_size = false,
int* ptr_num_valid_outputs = nullptr) {
const int output_size = max_output_size.scalar<int>()();
OP_REQUIRES(context, output_size >= 0,
errors::InvalidArgument("output size must be non-negative"));

std::vector<T> scores_data(num_boxes);
std::copy_n(scores.flat<T>().data(), num_boxes, scores_data.begin());
Expand Down Expand Up @@ -768,6 +770,9 @@ class NonMaxSuppressionV4Op : public OpKernel {
context, scores, num_boxes, max_output_size, iou_threshold_val,
score_threshold_val, dummy_soft_nms_sigma, similarity_fn,
return_scores_tensor_, pad_to_max_output_size_, &num_valid_outputs);
if (!context->status().ok()) {
return;
}

// Allocate scalar output tensor for number of indices computed.
Tensor* num_outputs_t = nullptr;
Expand Down Expand Up @@ -845,6 +850,9 @@ class NonMaxSuppressionV5Op : public OpKernel {
context, scores, num_boxes, max_output_size, iou_threshold_val,
score_threshold_val, soft_nms_sigma_val, similarity_fn,
return_scores_tensor_, pad_to_max_output_size_, &num_valid_outputs);
if (!context->status().ok()) {
return;
}

// Allocate scalar output tensor for number of indices computed.
Tensor* num_outputs_t = nullptr;
Expand Down

0 comments on commit 3a73627

Please sign in to comment.