diff --git a/mediapipe/tasks/c/components/containers/category_converter.cc b/mediapipe/tasks/c/components/containers/category_converter.cc index a38ef10b0f..b04f86c9a4 100644 --- a/mediapipe/tasks/c/components/containers/category_converter.cc +++ b/mediapipe/tasks/c/components/containers/category_converter.cc @@ -36,7 +36,9 @@ void CppConvertToCategory( void CppCloseCategory(Category* in) { free(in->category_name); + in->category_name = nullptr; free(in->display_name); + in->display_name = nullptr; } } // namespace mediapipe::tasks::c::components::containers diff --git a/mediapipe/tasks/c/components/containers/classification_result_converter.cc b/mediapipe/tasks/c/components/containers/classification_result_converter.cc index 64e62c3097..b7190c932e 100644 --- a/mediapipe/tasks/c/components/containers/classification_result_converter.cc +++ b/mediapipe/tasks/c/components/containers/classification_result_converter.cc @@ -60,17 +60,20 @@ void CppConvertToClassificationResult( void CppCloseClassificationResult(ClassificationResult* in) { for (uint32_t i = 0; i < in->classifications_count; ++i) { - auto classification_in = in->classifications[i]; + auto& classification_in = in->classifications[i]; for (uint32_t j = 0; j < classification_in.categories_count; ++j) { CppCloseCategory(&classification_in.categories[j]); } delete[] classification_in.categories; + classification_in.categories = nullptr; free(classification_in.head_name); + classification_in.head_name = nullptr; } delete[] in->classifications; + in->classifications = nullptr; } } // namespace mediapipe::tasks::c::components::containers diff --git a/mediapipe/tasks/c/components/containers/classification_result_converter_test.cc b/mediapipe/tasks/c/components/containers/classification_result_converter_test.cc index 59eb53aaff..cecef3ccda 100644 --- a/mediapipe/tasks/c/components/containers/classification_result_converter_test.cc +++ b/mediapipe/tasks/c/components/containers/classification_result_converter_test.cc @@ -99,4 +99,20 @@ TEST(ClassificationResultConverterTest, CppCloseClassificationResult(&c_classification_result); } +TEST(ClassificationResultConverterTest, FreesMemory) { + mediapipe::tasks::components::containers::ClassificationResult + cpp_classification_result = { + /* classifications= */ {{{}, 0, "foo"}}, + /* timestamp_ms= */ 42, + }; + + ClassificationResult c_classification_result; + CppConvertToClassificationResult(cpp_classification_result, + &c_classification_result); + EXPECT_NE(c_classification_result.classifications, nullptr); + + CppCloseClassificationResult(&c_classification_result); + EXPECT_EQ(c_classification_result.classifications, nullptr); +} + } // namespace mediapipe::tasks::c::components::containers