in tensorflow_lite_support/cc/task/vision/image_segmenter.cc [333:432]
StatusOr<SegmentationResult> ImageSegmenter::Postprocess(
const std::vector<const TfLiteTensor*>& output_tensors,
const FrameBuffer& frame_buffer, const BoundingBox& /*roi*/) {
if (output_tensors.size() != 1) {
return CreateStatusWithPayload(
StatusCode::kInternal,
absl::StrFormat("Expected 1 output tensors, found %d",
output_tensors.size()));
}
const TfLiteTensor* output_tensor = output_tensors[0];
SegmentationResult result;
Segmentation* segmentation = result.add_segmentation();
*segmentation->mutable_colored_labels() = {colored_labels_.begin(),
colored_labels_.end()};
// The output tensor has orientation `frame_buffer.orientation()`, as it has
// been produced from the pre-processed frame.
FrameBuffer::Orientation tensor_orientation = frame_buffer.orientation();
// The output tensor always has size `output_width_ x output_height_`
FrameBuffer::Dimension tensor_dimension = {output_width_, output_height_};
// The masks to produce from the output tensor need to be re-oriented in the
// unrotated frame of reference coordinates system, i.e. kTopLeft.
FrameBuffer::Orientation mask_orientation =
FrameBuffer::Orientation::kTopLeft;
// They may thus have swapped dimensions compared to the tensor if the
// rotation is 90° or 270°.
FrameBuffer::Dimension mask_dimension(tensor_dimension);
if (RequireDimensionSwap(frame_buffer.orientation(),
FrameBuffer::Orientation::kTopLeft)) {
mask_dimension.Swap();
}
segmentation->set_width(mask_dimension.width);
segmentation->set_height(mask_dimension.height);
// XY coordinates in the tensor, to be computed from mask_x and mask_y below.
int tensor_x;
int tensor_y;
if (options_->output_type() == ImageSegmenterOptions::CATEGORY_MASK) {
auto* category_mask = segmentation->mutable_category_mask();
category_mask->resize(mask_dimension.width * mask_dimension.height);
int pixel_offset = 0;
for (int mask_y = 0; mask_y < mask_dimension.height; ++mask_y) {
for (int mask_x = 0; mask_x < mask_dimension.width; ++mask_x) {
// Compute the coordinates (tensor_x, tensor_y) in the tensor with
// tensor_orientation = frame_buffer.orientation() corresponding to the
// coordinates (mask_x, mask_y) in the mask being filled with
// mask_orientation = kTopLeft, i.e. the orientation of the unrotated
// frame of reference.
OrientCoordinates(/*from_x=*/mask_x,
/*from_y=*/mask_y,
/*from_orientation=*/mask_orientation,
/*to_orientation=*/tensor_orientation,
/*from_dimension=*/mask_dimension,
/*to_x=*/&tensor_x,
/*to_y=*/&tensor_y);
int class_index = 0;
float max_confidence = 0.0f;
for (int d = 0; d < output_depth_; ++d) {
ASSIGN_OR_RETURN(
const float confidence,
GetOutputConfidence(*output_tensor, tensor_x, tensor_y, d));
if (confidence > max_confidence) {
class_index = d;
max_confidence = confidence;
}
}
(*category_mask)[pixel_offset++] = static_cast<char>(class_index);
}
}
} else if (options_->output_type() ==
ImageSegmenterOptions::CONFIDENCE_MASK) {
auto* confidence_masks = segmentation->mutable_confidence_masks();
for (int d = 0; d < output_depth_; ++d) {
confidence_masks->add_confidence_mask();
}
for (int mask_y = 0; mask_y < segmentation->height(); ++mask_y) {
for (int mask_x = 0; mask_x < segmentation->width(); ++mask_x) {
// See above.
OrientCoordinates(/*from_x=*/mask_x,
/*from_y=*/mask_y,
/*from_orientation=*/mask_orientation,
/*to_orientation=*/tensor_orientation,
/*from_dimension=*/mask_dimension,
/*to_x=*/&tensor_x,
/*to_y=*/&tensor_y);
for (int d = 0; d < output_depth_; ++d) {
ASSIGN_OR_RETURN(
float confidence,
GetOutputConfidence(*output_tensor, tensor_x, tensor_y, d));
confidence_masks->mutable_confidence_mask(d)->add_value(confidence);
}
}
}
}
return result;
}