Internal change

PiperOrigin-RevId: 524059494
This commit is contained in:
MediaPipe Team 2023-04-13 11:41:04 -07:00 committed by Copybara-Service
parent 055accece8
commit c63b30dff5

View File

@ -84,14 +84,12 @@ void Sigmoid(absl::Span<const float> values,
// Linearly interpolate the value between v0 and v1. Assume 0 <= t <= 1. // Linearly interpolate the value between v0 and v1. Assume 0 <= t <= 1.
float LinearInterpolate(float v0, float v1, float t) { float LinearInterpolate(float v0, float v1, float t) {
DCHECK(t >= 0 && t <= 1);
return v0 + (v1 - v0) * t; return v0 + (v1 - v0) * t;
} }
// Bilinearly interpolate the value between 4 points. Assume 0 <= t0, t1 <= 1. // Bilinearly interpolate the value between 4 points. Assume 0 <= t0, t1 <= 1.
float BilinearInterpolate(float v00, float v10, float v01, float v11, float t0, float BilinearInterpolate(float v00, float v10, float v01, float v11, float t0,
float t1) { float t1) {
DCHECK(t0 >= 0 && t0 <= 1 && t1 >= 0 && t1 <= 1);
return LinearInterpolate(LinearInterpolate(v00, v10, t0), return LinearInterpolate(LinearInterpolate(v00, v10, t0),
LinearInterpolate(v01, v11, t0), t1); LinearInterpolate(v01, v11, t0), t1);
} }
@ -120,16 +118,22 @@ Image ProcessForCategoryMaskCpu(const Shape& input_shape,
cv::Mat category_mask_mat_view = cv::Mat category_mask_mat_view =
mediapipe::formats::MatView(image_frame_ptr.get()); mediapipe::formats::MatView(image_frame_ptr.get());
const int input_channels = input_shape.channels; const int input_channels = input_shape.channels;
category_mask_mat_view.forEach<uint8_t>( category_mask_mat_view.forEach<uint8_t>([&tensors_buffer, &input_shape,
[&tensors_buffer, &input_shape, &width_scale, &height_scale, &width_scale, &height_scale,
&input_channels, &options](uint8_t& pixel, const int position[]) { &input_channels,
&options](uint8_t& pixel,
const int position[]) {
std::vector<float> confidence_scores(input_channels); std::vector<float> confidence_scores(input_channels);
int y0 = static_cast<int>(std::floor(position[0] * height_scale)); int y0 =
int x0 = static_cast<int>(std::floor(position[1] * width_scale)); static_cast<int>(std::max(std::floor(position[0] * height_scale), 0.f));
int y1 = static_cast<int>(std::ceil(position[0] * height_scale)); int x0 =
int x1 = static_cast<int>(std::ceil(position[1] * width_scale)); static_cast<int>(std::max(std::floor(position[1] * width_scale), 0.f));
float t0 = position[0] * height_scale - y0; int y1 = static_cast<int>(std::min(std::ceil(position[0] * height_scale),
float t1 = position[1] * width_scale - x0; input_shape.height - 1.f));
int x1 = static_cast<int>(std ::min(std::ceil(position[1] * width_scale),
input_shape.width - 1.f));
float t0 = std::max(std::min(position[0] * height_scale - y0, 1.f), 0.f);
float t1 = std::max(std::min(position[1] * width_scale - x0, 1.f), 0.f);
for (int i = 0; i < input_channels; ++i) { for (int i = 0; i < input_channels; ++i) {
confidence_scores[i] = BilinearInterpolate( confidence_scores[i] = BilinearInterpolate(
GetTensorElement(input_shape, tensors_buffer, x0, y0, i), GetTensorElement(input_shape, tensors_buffer, x0, y0, i),
@ -154,8 +158,7 @@ Image ProcessForCategoryMaskCpu(const Shape& input_shape,
pixel = static_cast<uint8_t>(confidence_scores[0] > 0.5f); pixel = static_cast<uint8_t>(confidence_scores[0] > 0.5f);
} else { } else {
const int maximum_category_idx = const int maximum_category_idx =
std::max_element(confidence_scores.begin(), std::max_element(confidence_scores.begin(), confidence_scores.end()) -
confidence_scores.end()) -
confidence_scores.begin(); confidence_scores.begin();
pixel = maximum_category_idx; pixel = maximum_category_idx;
} }