如何在TensorFlow C++ label_image示例中使用tensorflow::ops::NonMaxSuppression去重预测框

Hey there! I’ve worked with the TensorFlow C++ API for object detection before, so let’s walk through how to add tensorflow::ops::NonMaxSuppression to your label_image-based project to eliminate duplicate bounding boxes. Here’s a step-by-step breakdown tailored to your 10-class use case:

1. First, Extract Detection Outputs #

Since you’ve adapted the label_image example for localization, your model should output two core tensors:

• Bounding boxes: Shaped [num_boxes, 4], usually formatted as [y_min, x_min, y_max, x_max] (normalized to the range [0,1] in most pre-trained detection models)
• Confidence scores: Shaped [num_boxes], matching each box’s prediction confidence

First, grab these tensors after running your model:

// Assume your model's output nodes are named "detection_boxes" and "detection_scores"
tensorflow::Tensor boxes_tensor;
tensorflow::Tensor scores_tensor;
std::vector<std::pair<std::string, tensorflow::Tensor>> outputs;

// Run the session to retrieve outputs
auto status = session->Run(
    inputs,  // Your existing input tensor from the label_image setup
    {"detection_boxes", "detection_scores"},
    {},
    &outputs
);

if (!status.ok()) {
    LOG(ERROR) << "Model execution failed: " << status;
    return;
}

// Assign outputs to our tensors
boxes_tensor = outputs[0].second;
scores_tensor = outputs[1].second;

2. Prepare NMS Parameters #

NonMaxSuppression relies on a few key parameters to control box retention:

• max_output_size: Maximum number of boxes to keep (adjust based on your 10-class needs)
• iou_threshold: IOU (Intersection Over Union) threshold — boxes overlapping more than this get suppressed
• score_threshold: Optional filter to ignore low-confidence boxes upfront

Convert these settings to TensorFlow tensors:

// Configure your NMS rules
const int max_boxes_to_retain = 10;
const float iou_threshold = 0.5f;  // A common starting value; tweak for your data
const float score_cutoff = 0.3f;   // Ignore boxes with confidence < 30%

// Wrap parameters in TensorFlow tensors
tensorflow::Tensor max_output_tensor(tensorflow::DT_INT32, tensorflow::TensorShape());
max_output_tensor.scalar<int>()() = max_boxes_to_retain;

tensorflow::Tensor iou_tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape());
iou_tensor.scalar<float>()() = iou_threshold;

tensorflow::Tensor score_tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape());
score_tensor.scalar<float>()() = score_cutoff;

3. Execute the NonMaxSuppression Operation #

Build and run the NMS op using your existing session to avoid wasting resources:

// Create a scope for the NMS operation
tensorflow::Scope root = tensorflow::Scope::NewRootScope();

// Define the NMS op
auto selected_indices_op = tensorflow::ops::NonMaxSuppression(
    root.WithOpName("nms"),
    boxes_tensor,
    scores_tensor,
    max_output_tensor,
    tensorflow::ops::NonMaxSuppression::IouThreshold(iou_tensor),
    tensorflow::ops::NonMaxSuppression::ScoreThreshold(score_tensor)
);

// Run the op to get indices of boxes to keep
tensorflow::Tensor indices_tensor;
auto nms_status = session->Run(
    {},
    {selected_indices_op.name()},
    {},
    &indices_tensor
);

if (!nms_status.ok()) {
    LOG(ERROR) << "NMS execution failed: " << nms_status;
    return;
}

4. Extract Final Boxes and Scores #

The NMS op returns indices of the retained boxes. Use these to filter your original boxes and scores:

// Convert indices tensor to a usable integer vector
auto indices_flat = indices_tensor.flat<int>();
std::vector<int> kept_indices(indices_flat.size());
for (size_t i = 0; i < indices_flat.size(); ++i) {
    kept_indices[i] = indices_flat(i);
}

// Extract the final filtered boxes and scores
auto boxes_flat = boxes_tensor.flat_outer_dims<float>();
auto scores_flat = scores_tensor.flat<float>();

std::vector<std::array<float, 4>> final_boxes;
std::vector<float> final_scores;

for (int idx : kept_indices) {
    // Retrieve box coordinates (y_min, x_min, y_max, x_max)
    std::array<float, 4> box = {
        boxes_flat(idx, 0), boxes_flat(idx, 1),
        boxes_flat(idx, 2), boxes_flat(idx, 3)
    };
    final_boxes.push_back(box);
    final_scores.push_back(scores_flat(idx));
}

Critical Tips for Your 10-Class Project #

• Per-Class NMS: If your model outputs class labels (via a detection_classes tensor), run NMS separately for each class. Group boxes/scores by their class ID first, then apply NMS to each group — this prevents suppressing boxes from different classes that overlap.
• Coordinate Conversion: To draw boxes on the original image, multiply normalized coordinates by the image’s height and width to get pixel values.
• Tune Thresholds: Adjust iou_threshold and score_threshold based on test results — higher IOU thresholds keep more overlapping boxes, while higher score thresholds filter out more low-confidence predictions.

内容的提问来源于stack exchange，提问作者Android developer