Use box_output_format to decide keypoint order

PiperOrigin-RevId: 510464960
This commit is contained in:
MediaPipe Team 2023-02-17 10:28:08 -08:00 committed by Copybara-Service
parent 223c504d81
commit 1a60a0e2d6
2 changed files with 18 additions and 8 deletions

View File

@ -796,11 +796,19 @@ absl::Status TensorsToDetectionsCalculator::DecodeBoxes(
const int offset = i * num_coords_ + options_.keypoint_coord_offset() +
k * options_.num_values_per_keypoint();
float keypoint_y = raw_boxes[offset];
float keypoint_x = raw_boxes[offset + 1];
if (options_.reverse_output_order()) {
keypoint_x = raw_boxes[offset];
keypoint_y = raw_boxes[offset + 1];
float keypoint_y = 0.0;
float keypoint_x = 0.0;
switch (box_output_format_) {
case mediapipe::TensorsToDetectionsCalculatorOptions::UNSPECIFIED:
case mediapipe::TensorsToDetectionsCalculatorOptions::YXHW:
keypoint_y = raw_boxes[offset];
keypoint_x = raw_boxes[offset + 1];
break;
case mediapipe::TensorsToDetectionsCalculatorOptions::XYWH:
case mediapipe::TensorsToDetectionsCalculatorOptions::XYXY:
keypoint_x = raw_boxes[offset];
keypoint_y = raw_boxes[offset + 1];
break;
}
(*boxes)[offset] = keypoint_x / options_.x_scale() * anchors[i].w() +

View File

@ -116,14 +116,16 @@ message TensorsToDetectionsCalculatorOptions {
// Tells the calculator how to convert the detector output to bounding boxes.
// Replaces `reverse_output_order` to support more bbox output formats.
// As with `reverse_output_order`, this also informs calculator the order
// of keypoint predictions.
enum BoxFormat {
// if UNSPECIFIED, the calculator assumes YXHW
UNSPECIFIED = 0;
// [y_center, x_center, height, width]
// bbox [y_center, x_center, height, width], keypoint [y, x]
YXHW = 1;
// [x_center, y_center, width, height]
// bbox [x_center, y_center, width, height], keypoint [x, y]
XYWH = 2;
// [xmin, ymin, xmax, ymax]
// bbox [xmin, ymin, xmax, ymax], keypoint [x, y]
XYXY = 3;
}
optional BoxFormat box_format = 24 [default = UNSPECIFIED];