diff --git a/tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.mm b/tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.mm index f367fb32e..9eb82c7e3 100644 --- a/tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.mm +++ b/tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.mm @@ -108,14 +108,20 @@ - (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image error:(NSErro return nil; } - std::unique_ptr cppFrameBuffer = [image cppFrameBufferWithError:error]; + uint8_t *buffer = nil; + std::unique_ptr cppFrameBuffer = [image cppFrameBufferWithUnderlyingBuffer:&buffer + error:error]; if (!cppFrameBuffer) { + free(buffer); return nil; } StatusOr cppSearchResultStatus = _cppImageSearcher->Search(*cppFrameBuffer); + // Free the underlying buffer + free(buffer); + return [TFLSearchResult searchResultWithCppResult:cppSearchResultStatus error:error]; } @@ -129,9 +135,12 @@ - (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image return nil; } - std::unique_ptr cppFrameBuffer = [image cppFrameBufferWithError:error]; + uint8_t *buffer = nil; + std::unique_ptr cppFrameBuffer = [image cppFrameBufferWithUnderlyingBuffer:&buffer + error:error]; if (!cppFrameBuffer) { + free(buffer); return nil; } @@ -144,6 +153,9 @@ - (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image StatusOr cppSearchResultStatus = _cppImageSearcher->Search(*cppFrameBuffer, regionOfInterest); + // Free the underlying buffer + free(buffer); + return [TFLSearchResult searchResultWithCppResult:cppSearchResultStatus error:error]; } @end diff --git a/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.h b/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.h index f5b38903a..b33f814bc 100644 --- a/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.h +++ b/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.h @@ -31,14 +31,18 @@ NS_ASSUME_NONNULL_BEGIN * tflite::task::vision::FrameBuffer is used by the TFLite Task Vision C++ * library to hold the backing buffer of any image. * + * @param buffer Pointer to the memory location where underlying pixel buffer + * of the image should be saved. + * * @param error Pointer to the memory location where errors if any should be * saved. If @c NULL, no error will be saved. * * @return The FrameBuffer created from the gmlImage which can be used with the * TF Lite Task Vision C++ library. @c NULL in case of an error. */ -- (std::unique_ptr)cppFrameBufferWithError: - (NSError *_Nullable *)error; +- (std::unique_ptr) + cppFrameBufferWithUnderlyingBuffer:(uint8_t **)buffer + error:(NSError *_Nullable *)error; @end diff --git a/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.mm b/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.mm index 123e8d389..b25ab5aa7 100644 --- a/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.mm +++ b/tensorflow_lite_support/ios/task/vision/utils/sources/GMLImage+CppUtils.mm @@ -26,8 +26,10 @@ @implementation GMLImage (CppUtils) -- (std::unique_ptr)cppFrameBufferWithError:(NSError **)error { - uint8_t *buffer = [self bufferWithError:error]; +- (std::unique_ptr) + cppFrameBufferWithUnderlyingBuffer:(uint8_t **)buffer + error:(NSError *_Nullable *)error { + *buffer = [self bufferWithError:error]; if (!buffer) { return NULL; @@ -37,7 +39,7 @@ @implementation GMLImage (CppUtils) FrameBufferCpp::Format frame_buffer_format = FrameBufferCpp::Format::kRGB; StatusOr> frameBuffer = - CreateFromRawBuffer(buffer, {(int)bitmapSize.width, (int)bitmapSize.height}, + CreateFromRawBuffer(*buffer, {(int)bitmapSize.width, (int)bitmapSize.height}, frame_buffer_format, FrameBufferCpp::Orientation::kTopLeft); if (![TFLCommonCppUtils checkCppError:frameBuffer.status() toError:error]) {