Fix memory access issue in C layer

PiperOrigin-RevId: 569279959
This commit is contained in:
Sebastian Schmidt 2023-09-28 13:48:07 -07:00 committed by Copybara-Service
parent e169849041
commit 636cf99a3e
2 changed files with 9 additions and 6 deletions

View File

@ -25,10 +25,11 @@ void CppConvertToCategory(
Category* out) { Category* out) {
out->index = in.index; out->index = in.index;
out->score = in.score; out->score = in.score;
out->category_name = out->category_name = in.category_name.has_value()
in.category_name.has_value() ? in.category_name->c_str() : nullptr; ? strdup(in.category_name->c_str())
: nullptr;
out->display_name = out->display_name =
in.display_name.has_value() ? in.display_name->c_str() : nullptr; in.display_name.has_value() ? strdup(in.display_name->c_str()) : nullptr;
} }
} // namespace mediapipe::tasks::c::components::containers } // namespace mediapipe::tasks::c::components::containers

View File

@ -15,6 +15,8 @@ limitations under the License.
#include "mediapipe/tasks/c/components/containers/classification_result_converter.h" #include "mediapipe/tasks/c/components/containers/classification_result_converter.h"
#include <cstdint>
#include "mediapipe/tasks/c/components/containers/category.h" #include "mediapipe/tasks/c/components/containers/category.h"
#include "mediapipe/tasks/c/components/containers/category_converter.h" #include "mediapipe/tasks/c/components/containers/category_converter.h"
#include "mediapipe/tasks/cc/components/containers/classification_result.h" #include "mediapipe/tasks/cc/components/containers/classification_result.h"
@ -32,14 +34,14 @@ void CppConvertToClassificationResult(
out->classifications_count = in.classifications.size(); out->classifications_count = in.classifications.size();
out->classifications = new Classifications[out->classifications_count]; out->classifications = new Classifications[out->classifications_count];
for (uint32_t i = 0; i <= out->classifications_count; ++i) { for (uint32_t i = 0; i < out->classifications_count; ++i) {
auto classification_in = in.classifications[i]; auto classification_in = in.classifications[i];
auto classification_out = out->classifications[i]; auto classification_out = out->classifications[i];
classification_out.categories_count = classification_in.categories.size(); classification_out.categories_count = classification_in.categories.size();
classification_out.categories = classification_out.categories =
new Category[classification_out.categories_count]; new Category[classification_out.categories_count];
for (uint32_t j = 0; j <= classification_out.categories_count; ++j) { for (uint32_t j = 0; j < classification_out.categories_count; ++j) {
CppConvertToCategory(classification_in.categories[j], CppConvertToCategory(classification_in.categories[j],
&(classification_out.categories[j])); &(classification_out.categories[j]));
} }
@ -47,7 +49,7 @@ void CppConvertToClassificationResult(
classification_out.head_index = classification_in.head_index; classification_out.head_index = classification_in.head_index;
classification_out.head_name = classification_out.head_name =
classification_in.head_name.has_value() classification_in.head_name.has_value()
? classification_in.head_name.value().c_str() ? strdup(classification_in.head_name->c_str())
: nullptr; : nullptr;
} }
} }