Fixed memory leak in allocation of image pixel buffer

Russ76 · Oct 26, 2022 · 3e5c2ea · 3e5c2ea
1 parent 1407175
commit 3e5c2ea
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 7 deletions.
diff --git 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<FrameBufferCpp> cppFrameBuffer = [image cppFrameBufferWithError:error];
+  uint8_t *buffer = nil;
+  std::unique_ptr<FrameBufferCpp> cppFrameBuffer = [image cppFrameBufferWithUnderlyingBuffer:&buffer
+                                                                                       error:error];
 
   if (!cppFrameBuffer) {
+    free(buffer);
     return nil;
   }
 
   StatusOr<SearchResultCpp> 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<FrameBufferCpp> cppFrameBuffer = [image cppFrameBufferWithError:error];
+  uint8_t *buffer = nil;
+  std::unique_ptr<FrameBufferCpp> cppFrameBuffer = [image cppFrameBufferWithUnderlyingBuffer:&buffer
+                                                                                       error:error];
 
   if (!cppFrameBuffer) {
+    free(buffer);
     return nil;
   }
 
@@ -144,6 +153,9 @@ - (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image
   StatusOr<SearchResultCpp> 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
@@ -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<tflite::task::vision::FrameBuffer>)cppFrameBufferWithError:
-    (NSError *_Nullable *)error;
+- (std::unique_ptr<tflite::task::vision::FrameBuffer>)
+    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
@@ -26,8 +26,10 @@
 
 @implementation GMLImage (CppUtils)
 
-- (std::unique_ptr<FrameBufferCpp>)cppFrameBufferWithError:(NSError **)error {
-  uint8_t *buffer = [self bufferWithError:error];
+- (std::unique_ptr<tflite::task::vision::FrameBuffer>)
+    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<std::unique_ptr<FrameBufferCpp>> 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]) {