added ability to use shared runners

swittk · Sep 23, 2021 · e768cba · e768cba
1 parent 1406d56
commit e768cba
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 7 deletions.
diff --git a/ios/MediapipeFacemesh.m b/ios/MediapipeFacemesh.m
@@ -3,11 +3,19 @@
 #import "ChunkBasedFaceMeshRunner.h"
 
 @implementation MediapipeFacemesh {
+    NSMutableDictionary <NSNumber *, ChunkBasedFaceMeshRunner *>*runners;
     ChunkBasedFaceMeshRunner *runner;
 }
 
 RCT_EXPORT_MODULE()
 
+-(id)init {
+    if(self = [super init]) {
+        runners = [NSMutableDictionary new];
+    }
+    return self;
+}
+
 -(BOOL)requiresMainQueueSetup {
     return YES;
 }
@@ -35,8 +43,20 @@ -(BOOL)requiresMainQueueSetup {
         reject(@"NO_FILE_PATH", @"No File Path specified", error);
         return;
     }
-    runner = [ChunkBasedFaceMeshRunner new];
-    [[[runner processFilesAtPaths:filePaths]
+    NSNumber *runnerID = argumentsDict[@"runner"];
+    ChunkBasedFaceMeshRunner *myrunner;
+    if(runnerID) {
+        myrunner = runners[runnerID];
+        if(!myrunner) {
+            myrunner = [ChunkBasedFaceMeshRunner new];
+            runners[runnerID] = myrunner;
+        }
+    }
+    else {
+        runner = [ChunkBasedFaceMeshRunner new];
+        myrunner = runner;
+    }
+    [[[myrunner processFilesAtPaths:filePaths]
       then:^id _Nullable(NSArray<NSArray<NSArray<NSArray<NSNumber *>*>*>*>* value) {
         NSLog(@"result from runner %@", value);
         resolve(value);
@@ -46,6 +66,30 @@ -(BOOL)requiresMainQueueSetup {
     }];
 }
 
+RCT_REMAP_METHOD(listRunners,
+                 listRunners:(nullable NSDictionary *)argumentsDict
+                 withResolver:(RCTPromiseResolveBlock)resolve
+                 withRejecter:(RCTPromiseRejectBlock)reject)
+{
+//    argumentsDict[@""];
+    resolve([runners allKeys]);
+}
+
+RCT_REMAP_METHOD(freeRunner,
+                 freeRunner:(nonnull NSDictionary *)argumentsDict
+                 withResolver:(RCTPromiseResolveBlock)resolve
+                 withRejecter:(RCTPromiseRejectBlock)reject)
+{
+    NSNumber *runnerID = argumentsDict[@"runner"];
+    if(!runnerID) {
+        NSError *error = [[NSError alloc] initWithDomain:@"SKRNFaceMesh" code:404 userInfo:nil];
+        reject(@"NO_RUNNER_ID", @"No runner ID given", error);
+        return;
+    }
+    [runners removeObjectForKey:runnerID];
+    resolve([runners allKeys]);
+}
+
 // This is for checking the crash that occurs whenever I allocate FaceMeshIOSLib (  `[FaceMeshIOSLib new]`  )
 // The crash ends up to be from the logging library, `glog`, from its Lock() method which calls SAFE_PTHREAD internally.
 // I suspect that somehow Mediapipe static library and React native are both obtaining the same locks and thus screwing up.
@@ -83,8 +127,20 @@ -(BOOL)requiresMainQueueSetup {
     for(NSString *b64 in base64Images) {
         [images addObject:[self decodeBase64ToImage:b64]];
     }
-    runner = [ChunkBasedFaceMeshRunner new];
-    [[[runner processImages:images]
+    NSNumber *runnerID = argumentsDict[@"runner"];
+    ChunkBasedFaceMeshRunner *myrunner;
+    if(runnerID) {
+        myrunner = runners[runnerID];
+        if(!myrunner) {
+            myrunner = [ChunkBasedFaceMeshRunner new];
+            runners[runnerID] = myrunner;
+        }
+    }
+    else {
+        runner = [ChunkBasedFaceMeshRunner new];
+        myrunner = runner;
+    }
+    [[[myrunner processImages:images]
       then:^id _Nullable(NSArray<NSArray<NSArray<NSArray<NSNumber *>*>*>*>* value) {
         NSLog(@"result from runner %@", value);
         resolve(value);

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "react-native-mediapipe-facemesh",
-  "version": "0.5.0",
+  "version": "0.5.1",
   "description": "Mediapipe Facemesh fplugin or React Native",
   "main": "lib/commonjs/index",
   "module": "lib/module/index",

diff --git a/src/index.tsx b/src/index.tsx
@@ -6,13 +6,16 @@ type MediapipeFacemeshType = {
    * @param args files: URLs of files to process
    * @returns a [frames, faces, points, 3 (point_x_y_z)] tensor (array)
    */
-  runFaceMeshWithFiles(args: { files: string[] }): Promise<number[][][][]>;
+  runFaceMeshWithFiles(args: { files: string[], runner?: number }): Promise<number[][][][]>;
   /**
    * @param base64Images Base 64 images of frames wanted to process
    * @returns a [frames, faces, points, 3 (point_x_y_z)] tensor (array)
    */
-  runFaceMeshWithBase64Images(args: { base64Images: string[] }): Promise<number[][][][]>;
+  runFaceMeshWithBase64Images(args: { base64Images: string[], runner?: number }): Promise<number[][][][]>;
 
+  listRunners(): Promise<number[]>;
+
+  freeRunner(args: {runner: number}): Promise<number[]>;
 
   /** Debug */
   tryJustAlloc(args: { lib?: boolean }): Promise<void>;