feat!: make image segmentation generic, general refactor

.cspell-wordlist.txt

@@ -111,4 +111,7 @@ logprob
 RNFS
 pogodin
 kesha
-antonov
+antonov
+rfdetr
+basemodule
+IMAGENET

apps/computer-vision/app/image_segmentation/index.tsx

@@ -2,9 +2,8 @@ import Spinner from '../../components/Spinner';
 import { BottomBar } from '../../components/BottomBar';
 import { getImage } from '../../utils';
 import {
-  useImageSegmentation,
   DEEPLAB_V3_RESNET50,
-  DeeplabLabel,
+  useImageSegmentation,
 } from 'react-native-executorch';
 import {
   Canvas,
@@ -44,16 +43,20 @@ const numberToColor: number[][] = [
 ];
 
 export default function ImageSegmentationScreen() {
-  const model = useImageSegmentation({ model: DEEPLAB_V3_RESNET50 });
   const { setGlobalGenerating } = useContext(GeneratingContext);
-  useEffect(() => {
-    setGlobalGenerating(model.isGenerating);
-  }, [model.isGenerating, setGlobalGenerating]);
+  const { isReady, isGenerating, downloadProgress, forward } =
+    useImageSegmentation({
+      model: DEEPLAB_V3_RESNET50,
+    });
   const [imageUri, setImageUri] = useState('');
   const [imageSize, setImageSize] = useState({ width: 0, height: 0 });
   const [segImage, setSegImage] = useState<SkImage | null>(null);
   const [canvasSize, setCanvasSize] = useState({ width: 0, height: 0 });
 
+  useEffect(() => {
+    setGlobalGenerating(isGenerating);
+  }, [isGenerating, setGlobalGenerating]);
+
   const handleCameraPress = async (isCamera: boolean) => {
     const image = await getImage(isCamera);
     if (!image?.uri) return;
@@ -69,12 +72,8 @@ export default function ImageSegmentationScreen() {
     if (!imageUri || imageSize.width === 0 || imageSize.height === 0) return;
     try {
       const { width, height } = imageSize;
-      const output = await model.forward(imageUri, [DeeplabLabel.ARGMAX]);
-      const argmax = output[DeeplabLabel.ARGMAX] || [];
-      const uniqueValues = new Set<number>();
-      for (let i = 0; i < argmax.length; i++) {
-        uniqueValues.add(argmax[i]);
-      }
+      const output = await forward(imageUri, [], true);
+      const argmax = output.ARGMAX || [];
       const pixels = new Uint8Array(width * height * 4);
 
       for (let row = 0; row < height; row++) {
@@ -105,11 +104,11 @@ export default function ImageSegmentationScreen() {
     }
   };
 
-  if (!model.isReady) {
+  if (!isReady) {
     return (
       <Spinner
-        visible={!model.isReady}
-        textContent={`Loading the model ${(model.downloadProgress * 100).toFixed(0)} %`}
+        visible={!isReady}
+        textContent={`Loading the model ${(downloadProgress * 100).toFixed(0)} %`}
       />
     );
   }

docs/docs/03-hooks/02-computer-vision/useImageSegmentation.md

@@ -21,12 +21,15 @@ import {
   DEEPLAB_V3_RESNET50,
 } from 'react-native-executorch';
 
-const model = useImageSegmentation({ model: DEEPLAB_V3_RESNET50 });
+const model = useImageSegmentation({
+  model: DEEPLAB_V3_RESNET50,
+});
 
 const imageUri = 'file::///Users/.../cute_cat.png';
 
 try {
-  const outputDict = await model.forward(imageUri);
+  const result = await model.forward(imageUri);
+  // result.ARGMAX is an Int32Array of per-pixel class indices
 } catch (error) {
   console.error(error);
 }
@@ -36,9 +39,13 @@ try {
 
 `useImageSegmentation` takes [`ImageSegmentationProps`](../../06-api-reference/interfaces/ImageSegmentationProps.md) that consists of:
 
-- `model` containing [`modelSource`](../../06-api-reference/interfaces/ImageSegmentationProps.md#modelsource).
+- `model` - An object containing:
+  - `modelName` - The name of a built-in model. See [`ModelSources`](../../06-api-reference/type-aliases/ModelSources.md) for the list of supported models.
+  - `modelSource` - The location of the model binary (a URL or a bundled resource).
 - An optional flag [`preventLoad`](../../06-api-reference/interfaces/ImageSegmentationProps.md#preventload) which prevents auto-loading of the model.
 
+The hook is generic over the model config — TypeScript automatically infers the correct label type based on the `modelName` you provide. No explicit generic parameter is needed.
+
 You need more details? Check the following resources:
 
 - For detailed information about `useImageSegmentation` arguments check this section: [`useImageSegmentation` arguments](../../06-api-reference/functions/useImageSegmentation.md#parameters).
@@ -47,45 +54,70 @@ You need more details? Check the following resources:
 
 ### Returns
 
-`useImageSegmentation` returns an object called `ImageSegmentationType` containing bunch of functions to interact with image segmentation models. To get more details please read: [`ImageSegmentationType` API Reference](../../06-api-reference/interfaces/ImageSegmentationType.md).
+`useImageSegmentation` returns an [`ImageSegmentationType`](../../06-api-reference/interfaces/ImageSegmentationType.md) object containing:
+
+- `isReady` - Whether the model is loaded and ready to process images.
+- `isGenerating` - Whether the model is currently processing an image.
+- `error` - An error object if the model failed to load or encountered a runtime error.
+- `downloadProgress` - A value between 0 and 1 representing the download progress of the model binary.
+- `forward` - A function to run inference on an image.
 
 ## Running the model
 
-To run the model, you can use the [`forward`](../../06-api-reference/interfaces/ImageSegmentationType.md#forward) method. It accepts three arguments: a required image - can be a remote URL, a local file URI, or a base64-encoded image (whole URI or only raw base64), an optional list of classes, and an optional flag whether to resize the output to the original dimensions.
+To run the model, use the [`forward`](../../06-api-reference/interfaces/ImageSegmentationType.md#forward) method. It accepts three arguments:
 
-- The image can be a remote URL, a local file URI, or a base64-encoded image.
-- The [`classesOfInterest`](../../06-api-reference/interfaces/ImageSegmentationType.md#classesofinterest) list contains classes for which to output the full results. By default the list is empty, and only the most probable classes are returned (essentially an arg max for each pixel). Look at [`DeeplabLabel`](../../06-api-reference/enumerations/DeeplabLabel.md) enum for possible classes.
-- The [`resizeToInput`](../../06-api-reference/interfaces/ImageSegmentationType.md#resizetoinput) flag specifies whether the output will be rescaled back to the size of the input image. The default is `true`. The model runs inference on a scaled (probably smaller) version of your image (224x224 for `DEEPLAB_V3_RESNET50`). If you choose to resize, the output will be `number[]` of size `width * height` of your original image.
+- [`imageSource`](../../06-api-reference/interfaces/ImageSegmentationType.md#forward) (required) - The image to segment. Can be a remote URL, a local file URI, or a base64-encoded image (whole URI or only raw base64).
+- [`classesOfInterest`](../../06-api-reference/interfaces/ImageSegmentationType.md#forward) (optional) - An array of label keys indicating which per-class probability masks to include in the output. Defaults to `[]` (no class masks). The `ARGMAX` map is always returned regardless of this parameter.
+- [`resizeToInput`](../../06-api-reference/interfaces/ImageSegmentationType.md#forward) (optional) - Whether to resize the output masks to the original input image dimensions. Defaults to `true`. If `false`, returns the raw model output dimensions (e.g. 224x224 for `DEEPLAB_V3_RESNET50`).
 
 :::warning
 Setting `resizeToInput` to `false` will make `forward` faster.
 :::
 
-[`forward`](../../06-api-reference/interfaces/ImageSegmentationType.md#forward) returns a promise which can resolve either to an error or a dictionary containing number arrays with size depending on [`resizeToInput`](../../06-api-reference/interfaces/ImageSegmentationType.md#resizetoinput):
+`forward` returns a promise resolving to an object containing:
+
+- `ARGMAX` - An `Int32Array` where each element is the class index with the highest probability for that pixel.
+- For each label included in `classesOfInterest`, a `Float32Array` of per-pixel probabilities for that class.
 
-- For the key [`DeeplabLabel.ARGMAX`](../../06-api-reference/enumerations/DeeplabLabel.md#argmax) the array contains for each pixel an integer corresponding to the class with the highest probability.
-- For every other key from [`DeeplabLabel`](../../06-api-reference/enumerations/DeeplabLabel.md), if the label was included in [`classesOfInterest`](../../06-api-reference/interfaces/ImageSegmentationType.md#classesofinterest) the dictionary will contain an array of floats corresponding to the probability of this class for every pixel.
+The return type is fully typed — TypeScript narrows it based on the labels you pass in `classesOfInterest`.
 
 ## Example
 
 ```typescript
+import {
+  useImageSegmentation,
+  DEEPLAB_V3_RESNET50,
+  DeeplabLabel,
+} from 'react-native-executorch';
+
 function App() {
-  const model = useImageSegmentation({ model: DEEPLAB_V3_RESNET50 });
+  const model = useImageSegmentation({
+    model: DEEPLAB_V3_RESNET50,
+  });
 
-  // ...
-  const imageUri = 'file::///Users/.../cute_cat.png';
+  const handleSegment = async () => {
+    if (!model.isReady) return;
+
+    const imageUri = 'file::///Users/.../cute_cat.png';
+
+    try {
+      const result = await model.forward(imageUri, ['CAT', 'PERSON'], true);
+
+      // result.ARGMAX — Int32Array of per-pixel class indices
+      // result.CAT — Float32Array of per-pixel probabilities for CAT
+      // result.PERSON — Float32Array of per-pixel probabilities for PERSON
+    } catch (error) {
+      console.error(error);
+    }
+  };
 
-  try {
-    const outputDict = await model.forward(imageUri, [DeeplabLabel.CAT], true);
-  } catch (error) {
-    console.error(error);
-  }
   // ...
 }
 ```
 
 ## Supported models
 
-| Model                                                                                            | Number of classes | Class list                                                          |
-| ------------------------------------------------------------------------------------------------ | ----------------- | ------------------------------------------------------------------- |
-| [deeplabv3_resnet50](https://huggingface.co/software-mansion/react-native-executorch-deeplab-v3) | 21                | [DeeplabLabel](../../06-api-reference/enumerations/DeeplabLabel.md) |
+| Model                                                                                            | Number of classes | Class list                                                                                |
+| ------------------------------------------------------------------------------------------------ | ----------------- | ----------------------------------------------------------------------------------------- |
+| [deeplabv3_resnet50](https://huggingface.co/software-mansion/react-native-executorch-deeplab-v3) | 21                | [DeeplabLabel](../../06-api-reference/enumerations/DeeplabLabel.md)                       |
+| selfie-segmentation                                                                              | 2                 | [SelfieSegmentationLabel](../../06-api-reference/enumerations/SelfieSegmentationLabel.md) |

docs/docs/04-typescript-api/02-computer-vision/ImageSegmentationModule.md

@@ -19,14 +19,15 @@ import {
 
 const imageUri = 'path/to/image.png';
 
-// Creating an instance
-const imageSegmentationModule = new ImageSegmentationModule();
-
-// Loading the model
-await imageSegmentationModule.load(DEEPLAB_V3_RESNET50);
+// Creating an instance from a built-in model
+const segmentation = await ImageSegmentationModule.fromModelName({
+  modelName: 'deeplab-v3',
+  modelSource: DEEPLAB_V3_RESNET50,
+});
 
 // Running the model
-const outputDict = await imageSegmentationModule.forward(imageUri);
+const result = await segmentation.forward(imageUri);
+// result.ARGMAX — Int32Array of per-pixel class indices
 ```
 
 ### Methods
@@ -35,34 +36,75 @@ All methods of `ImageSegmentationModule` are explained in details here: [`ImageS
 
 ## Loading the model
 
-To initialize the module, create an instance and call the [`load`](../../06-api-reference/classes/ImageSegmentationModule.md#load) method with the following parameters:
+`ImageSegmentationModule` uses static factory methods instead of `new()` + `load()`. There are two ways to create an instance:
+
+### Built-in models — `fromModelName`
+
+Use [`fromModelName`](../../06-api-reference/classes/ImageSegmentationModule.md#frommodelname) for models that ship with built-in label maps and preprocessing configs:
+
+```typescript
+const segmentation = await ImageSegmentationModule.fromModelName(
+  DEEPLAB_V3_RESNET50,
+  (progress) => console.log(`Download: ${Math.round(progress * 100)}%`)
+);
+```
 
-- [`model`](../../06-api-reference/classes/ImageSegmentationModule.md#model) - Object containing:
-  - [`modelSource`](../../06-api-reference/classes/ImageSegmentationModule.md#modelsource) - Location of the used model.
+The `config` parameter is a discriminated union — TypeScript ensures you provide the correct fields for each model name. Available built-in models: `'deeplab-v3'`, `'selfie-segmentation'`.
 
-- [`onDownloadProgressCallback`](../../06-api-reference/classes/ImageSegmentationModule.md#ondownloadprogresscallback) - Callback to track download progress.
+### Custom models — `fromCustomConfig`
 
-This method returns a promise, which can resolve to an error or void.
+Use [`fromCustomConfig`](../../06-api-reference/classes/ImageSegmentationModule.md#fromcustomconfig) for custom-exported segmentation models with your own label map:
+
+```typescript
+const MyLabels = { BACKGROUND: 0, FOREGROUND: 1 } as const;
+
+const segmentation = await ImageSegmentationModule.fromCustomConfig(
+  'https://example.com/custom_model.pte',
+  {
+    labelMap: MyLabels,
+    preprocessorConfig: {
+      normMean: [0.485, 0.456, 0.406],
+      normStd: [0.229, 0.224, 0.225],
+    },
+  }
+);
+```
+
+The `preprocessorConfig` is optional. If omitted, no input normalization is applied. The module instance will be typed to your custom label map — `forward` will accept and return keys from `MyLabels`.
 
 For more information on loading resources, take a look at [loading models](../../01-fundamentals/02-loading-models.md) page.
 
 ## Running the model
 
-To run the model, you can use the [`forward`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) method on the module object. It accepts three arguments: a required image - can be a remote URL, a local file URI, or a base64-encoded image (whole URI or only raw base64), an optional list of classes, and an optional flag whether to resize the output to the original dimensions.
+To run the model, use the [`forward`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) method. It accepts three arguments:
 
-- The image can be a remote URL, a local file URI, or a base64-encoded image.
-- The [`classesOfInterest`](../../06-api-reference/classes/ImageSegmentationModule.md#classesofinterest) list contains classes for which to output the full results. By default the list is empty, and only the most probable classes are returned (essentially an arg max for each pixel). Look at [`DeeplabLabel`](../../06-api-reference/enumerations/DeeplabLabel.md) enum for possible classes.
-- The [`resizeToInput`](../../06-api-reference/classes/ImageSegmentationModule.md#resizetoinput) flag specifies whether the output will be rescaled back to the size of the input image. The default is `true`. The model runs inference on a scaled (probably smaller) version of your image (224x224 for the `DEEPLAB_V3_RESNET50`). If you choose to resize, the output will be `number[]` of size `width * height` of your original image.
+- [`imageSource`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) (required) - The image to segment. Can be a remote URL, a local file URI, or a base64-encoded image (whole URI or only raw base64).
+- [`classesOfInterest`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) (optional) - An array of label keys indicating which per-class probability masks to include in the output. Defaults to `[]`. The `ARGMAX` map is always returned regardless.
+- [`resizeToInput`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) (optional) - Whether to resize the output masks to the original input image dimensions. Defaults to `true`. If `false`, returns the raw model output dimensions.
 
 :::warning
-Setting `resize` to true will make `forward` slower.
+Setting `resizeToInput` to `false` will make `forward` faster.
 :::
 
-[`forward`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) returns a promise which can resolve either to an error or a dictionary containing number arrays with size depending on [`resizeToInput`](../../06-api-reference/classes/ImageSegmentationModule.md#resizetoinput):
+`forward` returns a promise resolving to an object containing:
+
+- `ARGMAX` - An `Int32Array` where each element is the class index with the highest probability for that pixel.
+- For each label included in `classesOfInterest`, a `Float32Array` of per-pixel probabilities for that class.
 
-- For the key [`DeeplabLabel.ARGMAX`](../../06-api-reference/enumerations/DeeplabLabel.md#argmax) the array contains for each pixel an integer corresponding to the class with the highest probability.
-- For every other key from [`DeeplabLabel`](../../06-api-reference/enumerations/DeeplabLabel.md), if the label was included in [`classesOfInterest`](../../06-api-reference/classes/ImageSegmentationModule.md#classesofinterest) the dictionary will contain an array of floats corresponding to the probability of this class for every pixel.
+The return type narrows based on the labels passed in `classesOfInterest`:
+
+```typescript
+// Only ARGMAX in the result
+const result = await segmentation.forward(imageUri);
+result.ARGMAX; // Int32Array
+
+// ARGMAX + requested class masks
+const result = await segmentation.forward(imageUri, ['CAT', 'DOG']);
+result.ARGMAX; // Int32Array
+result.CAT; // Float32Array
+result.DOG; // Float32Array
+```
 
 ## Managing memory
 
-The module is a regular JavaScript object, and as such its lifespan will be managed by the garbage collector. In most cases this should be enough, and you should not worry about freeing the memory of the module yourself, but in some cases you may want to release the memory occupied by the module before the garbage collector steps in. In this case use the method [`delete`](../../06-api-reference/classes/ImageSegmentationModule.md#delete) on the module object you will no longer use, and want to remove from the memory. Note that you cannot use [`forward`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) after [`delete`](../../06-api-reference/classes/ImageSegmentationModule.md#delete) unless you load the module again.
+The module is a regular JavaScript object, and as such its lifespan will be managed by the garbage collector. In most cases this should be enough, and you should not worry about freeing the memory of the module yourself, but in some cases you may want to release the memory occupied by the module before the garbage collector steps in. In this case use the method [`delete`](../../06-api-reference/classes/ImageSegmentationModule.md#delete) on the module object you will no longer use, and want to remove from the memory. Note that you cannot use [`forward`](../../06-api-reference/classes/ImageSegmentationModule.md#forward) after [`delete`](../../06-api-reference/classes/ImageSegmentationModule.md#delete) unless you create a new instance.