Merge pull request #891 from khanhlvg:ios-image-searcher-api-impl

PiperOrigin-RevId: 482325865

Author: tflite-support-robot

tensorflow_lite_support/ios/BUILD

@@ -153,6 +153,10 @@ objc_library(
     hdrs = [
         "sources/TFLCommonCppUtils.h",
     ],
+    copts = [
+        "-ObjC++",
+        "-std=c++17",
+    ],
     module_name = "TFLCommonCppUtils",
     deps = [
         ":TFLCommon",

tensorflow_lite_support/ios/task/processor/sources/TFLSearchOptions.m

@@ -22,6 +22,7 @@ - (instancetype)init {
     // maxResults will be 0 at the time of initialization. Setting it to 5 since max_results
     // defaults to 5 in search_options.proto.
     _maxResults = 5;
+    _indexFile = [[TFLExternalFile alloc] init];
   }
   return self;
 }

tensorflow_lite_support/ios/task/vision/BUILD

@@ -89,6 +89,7 @@ objc_library(
     module_name = "TFLImageSearcher",
     deps = [
         "//tensorflow_lite_support/cc/task/vision:image_searcher",
+        "//tensorflow_lite_support/ios:TFLCommonCppUtils",
         "//tensorflow_lite_support/ios:TFLCommonUtils",
         "//tensorflow_lite_support/ios/task/core:TFLBaseOptionsCppHelpers",
         "//tensorflow_lite_support/ios/task/processor:TFLEmbeddingOptionsHelpers",

tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.h

@@ -96,15 +96,41 @@ NS_SWIFT_NAME(ImageSearcher)
  *    camera and get the frames for inference, you must request for this format
  *    from AVCaptureVideoDataOutput. Otherwise your inference results will be wrong.
  *
- * @param image An image on which embedding extraction is to be performed, followed by
- * nearest-neighbor search in the index, represented as a `GMLImage`.
- *
+ * @param image An image on which to perform embedding extraction, followed by a nearest-neighbor
+ * search in the index, represented as a `GMLImage`.
+
  * @return A `TFLSearchResult`. `nil` if there is an error encountered during embedding extraction
  * and nearest neighbor search. Please see `TFLSearchResult` for more details.
  */
 - (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image
                                            error:(NSError **)error NS_SWIFT_NAME(search(mlImage:));
 
+/**
+ * Performs embedding extraction on the given GMLImage, followed by nearest-neighbor search in the
+ * index on the pixels within the specified region of interest of the given
+ * `GMLImage`.
+ *
+ * @discussion This method currently supports inference on only following type of images:
+ * 1. RGB and RGBA images for `GMLImageSourceTypeImage`.
+ * 2. kCVPixelFormatType_32BGRA for `GMLImageSourceTypePixelBuffer` and
+ *    `GMLImageSourceTypeSampleBuffer`. If you are using `AVCaptureSession` to setup
+ *    camera and get the frames for inference, you must request for this format
+ *    from AVCaptureVideoDataOutput. Otherwise your classification
+ *    results will be wrong.
+ *
+ * @param image An image on which to perform embedding extraction, followed by a nearest-neighbor
+ * search in the index, represented as a `GMLImage`.
+ *
+ * @param roi A CGRect specifying the region of interest within the given `GMLImage`.
+ *
+ * @return A TFLClassificationResult with one set of results per image classifier head. `nil` if
+ * there is an error encountered during classification.
+ */
+- (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image
+                                regionOfInterest:(CGRect)roi
+                                           error:(NSError **)error
+    NS_SWIFT_NAME(search(mlImage:regionOfInterest:));
+
 - (instancetype)init NS_UNAVAILABLE;
 
 @end

tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.mm

@@ -14,7 +14,7 @@
  ==============================================================================*/
 #import "tensorflow_lite_support/ios/task/vision/sources/TFLImageSearcher.h"
 #import "tensorflow_lite_support/ios/sources/TFLCommon.h"
-#import "tensorflow_lite_support/ios/sources/TFLCommonUtils.h"
+#import "tensorflow_lite_support/ios/sources/TFLCommonCppUtils.h"
 #import "tensorflow_lite_support/ios/task/core/sources/TFLBaseOptions+CppHelpers.h"
 #import "tensorflow_lite_support/ios/task/processor/sources/TFLEmbeddingOptions+Helpers.h"
 #import "tensorflow_lite_support/ios/task/processor/sources/TFLSearchOptions+Helpers.h"
@@ -71,16 +71,17 @@ - (ImageSearcherOptionsCpp)cppOptions {
 
 @implementation TFLImageSearcher
 
-- (nullable instancetype)initWithCppImageSearcherOptions:(ImageSearcherOptionsCpp)cppOptions {
+- (nullable instancetype)initWithCppImageSearcherOptions:(ImageSearcherOptionsCpp)cppOptions
+                                                   error:(NSError **)error {
   self = [super init];
   if (self) {
     StatusOr<std::unique_ptr<ImageSearcherCpp>> cppImageSearcher =
         ImageSearcherCpp::CreateFromOptions(cppOptions);
-    if (cppImageSearcher.ok()) {
-      _cppImageSearcher = std::move(cppImageSearcher.value());
-    } else {
+    if (![TFLCommonCppUtils checkCppError:cppImageSearcher.status() toError:error]) {
       return nil;
     }
+
+    _cppImageSearcher = std::move(cppImageSearcher.value());
   }
   return self;
 }
@@ -96,7 +97,7 @@ + (nullable instancetype)imageSearcherWithOptions:(TFLImageSearcherOptions *)opt
 
   ImageSearcherOptionsCpp cppOptions = [options cppOptions];
 
-  return [[TFLImageSearcher alloc] initWithCppImageSearcherOptions:cppOptions];
+  return [[TFLImageSearcher alloc] initWithCppImageSearcherOptions:cppOptions error:error];
 }
 
 - (nullable TFLSearchResult *)searchWithGMLImage:(GMLImage *)image error:(NSError **)error {

tensorflow_lite_support/ios/test/task/vision/image_searcher/BUILD

@@ -13,10 +13,12 @@ objc_library(
     srcs = ["TFLImageSearcherTests.m"],
     data = [
         "//tensorflow_lite_support/cc/test/testdata/task/vision:test_images",
+        "//tensorflow_lite_support/cc/test/testdata/task/vision:test_indices",
         "//tensorflow_lite_support/cc/test/testdata/task/vision:test_models",
     ],
     tags = TFL_DEFAULT_TAGS,
     deps = [
+        "//tensorflow_lite_support/ios:TFLCommon",
         "//tensorflow_lite_support/ios/task/vision:TFLImageSearcher",
         "//tensorflow_lite_support/ios/task/vision/utils:GMLImageUtils",
     ],