Add a way to filter what segmentation classes are returned

Author: JakubGonera

ios/RnExecutorch/ImageSegmentation.mm

@@ -34,12 +34,13 @@ - (void)loadModule:(NSString *)modelSource
 }
 
 - (void)forward:(NSString *)input
+        classesOfInterest:(NSArray *)classesOfInterest
         resolve:(RCTPromiseResolveBlock)resolve
          reject:(RCTPromiseRejectBlock)reject {
 
   @try {
     cv::Mat image = [ImageProcessor readImage:input];
-    NSDictionary *result= [model runModel:image];
+    NSDictionary *result = [model runModel:image returnClasses:classesOfInterest];
 
     resolve(result);
     return;

ios/RnExecutorch/models/image_segmentation/ImageSegmentationModel.h

@@ -3,7 +3,8 @@
 
 @interface ImageSegmentationModel : BaseModel
 - (cv::Size)getModelImageSize;
-- (NSDictionary *)runModel:(cv::Mat &)input;
+- (NSDictionary *)runModel:(cv::Mat &)input
+                  returnClasses:(NSArray *)classesOfInterest;
 
 @end
 

ios/RnExecutorch/models/image_segmentation/ImageSegmentationModel.mm

@@ -1,12 +1,14 @@
 #import "ImageSegmentationModel.h"
+#import <unordered_set>
 #import "../../utils/ImageProcessor.h"
 #import "../../utils/Numerical.h"
 #import "opencv2/opencv.hpp"
 #import "Constants.h"
 
 @interface ImageSegmentationModel ()
   - (NSArray *)preprocess:(cv::Mat &)input;
-  - (NSDictionary *)postprocess:(NSArray *)output;
+  - (NSDictionary *)postprocess:(NSArray *)output
+                    returnClasses:(NSArray *)classesOfInterest;
 @end
 
 @implementation ImageSegmentationModel {
@@ -35,7 +37,8 @@ - (NSArray *)preprocess:(cv::Mat &)input {
   return modelInput;
 }
 
-- (NSDictionary *)postprocess:(NSArray *)output {
+- (NSDictionary *)postprocess:(NSArray *)output
+                  returnClasses:(NSArray *)classesOfInterest{
   cv::Size modelImageSize = [self getModelImageSize];
 
   std::size_t numLabels = deeplabv3_resnet50_labels.size();
@@ -87,25 +90,34 @@ - (NSDictionary *)postprocess:(NSArray *)output {
     maxArg.at<int>(row, col) = maxArgIndex;
   }
 
+  std::unordered_set<std::string> labelSet;
+
+  for (id label in classesOfInterest) {
+      labelSet.insert(std::string([label UTF8String]));
+  }
+
   NSMutableDictionary *result = [NSMutableDictionary dictionary];
-  
+
   // Convert to NSArray and populate the final dictionary
   for (std::size_t label = 0; label < numLabels; ++label) {
-    NSString *labelString = @(deeplabv3_resnet50_labels[label].c_str());
-    NSMutableArray *arr = matToNSArray<double>(resizedLabelScores[label]);
-    result[labelString] = arr;
+    if (labelSet.contains(deeplabv3_resnet50_labels[label])){
+        NSString *labelString = @(deeplabv3_resnet50_labels[label].c_str());
+        NSArray *arr = matToNSArray<double>(resizedLabelScores[label]);
+        result[labelString] = arr;
+    }
   }
 
   result[@"argmax"] = matToNSArray<int>(maxArg);
 
   return result;
 }
 
-- (NSDictionary *)runModel:(cv::Mat &)input {
+- (NSDictionary *)runModel:(cv::Mat &)input
+                  returnClasses:(NSArray *)classesOfInterest {
   NSArray *modelInput = [self preprocess:input];
   NSArray *result = [self forward:modelInput];
 
-  NSDictionary *output = [self postprocess:result[0]];
+  NSDictionary *output = [self postprocess:result[0] returnClasses:classesOfInterest];
 
   return output;
 }

src/hooks/computer_vision/useImageSegmentation.ts

@@ -1,6 +1,7 @@
-import { useState } from 'react';
+import { useState, useEffect } from 'react';
 import { _ImageSegmentationModule } from '../../native/RnExecutorchModules';
-import { useModule } from '../useModule';
+import { fetchResource } from '../../utils/fetchResource';
+import { ETError, getError } from '../../Error';
 
 interface Props {
   modelSource: string | number;
@@ -13,19 +14,52 @@ export const useImageSegmentation = ({
   isReady: boolean;
   isGenerating: boolean;
   downloadProgress: number;
-  forward: (input: string) => Promise<{ [category: string]: number[] }>;
+  forward: (
+    input: string,
+    classesOfInterest?: string[]
+  ) => Promise<{ [category: string]: number[] }>;
 } => {
   const [module, _] = useState(() => new _ImageSegmentationModule());
-  const {
-    error,
-    isReady,
-    isGenerating,
-    downloadProgress,
-    forwardImage: forward,
-  } = useModule({
-    modelSource,
-    module,
-  });
+  const [error, setError] = useState<null | string>(null);
+  const [isReady, setIsReady] = useState(false);
+  const [downloadProgress, setDownloadProgress] = useState(0);
+  const [isGenerating, setIsGenerating] = useState(false);
+
+  useEffect(() => {
+    const loadModel = async () => {
+      if (!modelSource) return;
+
+      try {
+        setIsReady(false);
+        const fileUri = await fetchResource(modelSource, setDownloadProgress);
+        await module.loadModule(fileUri);
+        setIsReady(true);
+      } catch (e) {
+        setError(getError(e));
+      }
+    };
+
+    loadModel();
+  }, [modelSource, module]);
+
+  const forward = async (input: string, classesOfInterest?: string[]) => {
+    if (!isReady) {
+      throw new Error(getError(ETError.ModuleNotLoaded));
+    }
+    if (isGenerating) {
+      throw new Error(getError(ETError.ModelGenerating));
+    }
+
+    try {
+      setIsGenerating(true);
+      const output = await module.forward(input, classesOfInterest || []);
+      return output;
+    } catch (e) {
+      throw new Error(getError(e));
+    } finally {
+      setIsGenerating(false);
+    }
+  };
 
   return { error, isReady, isGenerating, downloadProgress, forward };
 };

src/modules/computer_vision/BaseCVModule.ts

@@ -3,15 +3,13 @@ import {
   _StyleTransferModule,
   _ObjectDetectionModule,
   _ClassificationModule,
-  _ImageSegmentationModule,
 } from '../../native/RnExecutorchModules';
 import { getError } from '../../Error';
 
 export class BaseCVModule extends BaseModule {
   static module:
     | _StyleTransferModule
     | _ObjectDetectionModule
-    | _ImageSegmentationModule
     | _ClassificationModule;
 
   static async forward(input: string) {

src/modules/computer_vision/ImageSegmentationModule.ts

@@ -1,12 +1,19 @@
-import { BaseCVModule } from './BaseCVModule';
+import { BaseModule } from '../BaseModule';
 import { _ImageSegmentationModule } from '../../native/RnExecutorchModules';
+import { getError } from '../../Error';
 
-export class ImageSegmentationModule extends BaseCVModule {
+export class ImageSegmentationModule extends BaseModule {
   static module = new _ImageSegmentationModule();
 
-  static async forward(input: string) {
-    return await (super.forward(input) as ReturnType<
-      _ImageSegmentationModule['forward']
-    >);
+  static async forward(input: string, classesOfInteres?: string[]) {
+    console.log('# classes: ', classesOfInteres?.length);
+    try {
+      return await (this.module.forward(
+        input,
+        classesOfInteres || []
+      ) as ReturnType<_ImageSegmentationModule['forward']>);
+    } catch (e) {
+      throw new Error(getError(e));
+    }
   }
 }

src/native/NativeImageSegmentation.ts

@@ -4,7 +4,10 @@ import { TurboModuleRegistry } from 'react-native';
 export interface Spec extends TurboModule {
   loadModule(modelSource: string): Promise<number>;
 
-  forward(input: string): Promise<{ [category: string]: number[] }>;
+  forward(
+    input: string,
+    classesOfInterest: string[]
+  ): Promise<{ [category: string]: number[] }>;
 }
 
 export default TurboModuleRegistry.get<Spec>('ImageSegmentation');

src/native/RnExecutorchModules.ts

@@ -117,9 +117,10 @@ const OCR = OCRSpec
 
 class _ImageSegmentationModule {
   async forward(
-    input: string
+    input: string,
+    classesOfInteres: string[]
   ): ReturnType<ImageSegmentationInterface['forward']> {
-    return await ImageSegmentation.forward(input);
+    return await ImageSegmentation.forward(input, classesOfInteres);
   }
   async loadModule(
     modelSource: string | number