feat!: make image segmentation generic, general refactor (software-mansion#814)

## Description

Refactors image segmentation into a generic, multi-model architecture.
Previously the module was hardcoded to DeepLab V3 — now it supports
multiple built-in models (DeepLab V3, selfie segmentation, RF-DETR) and
custom user-provided models with type-safe label maps.

**Key changes:**

- **C++ base class**: Extracted `BaseImageSegmentation` with virtual
`preprocess()`/`postprocess()` methods. `ImageSegmentation` is now a
thin subclass. This allows future models to override preprocessing (e.g.
different normalization) or postprocessing without duplicating the
pipeline.
- **Optional normalization in C++**: `readImageToTensor` now accepts
optional `normMean`/`normStd` params, eliminating duplicated
normalization logic. **also, imo it would be a good idea to do such
factories for the entire API**
- **Generic TypeScript module**: `ImageSegmentationModule<T>` is generic
over model name or custom `LabelEnum`. Two static factories:
`fromModelName()` (built-in models with auto label resolution) and
`fromCustomConfig()` (custom models with user-provided labels).
- **Generic hook**: `useImageSegmentation` infers the model's label
types from the config — no explicit generic parameter needed.
`forward()` return type narrows based on `classesOfInterest` passed in.
- **Correct return types**: `forward()` now returns `Record<'ARGMAX',
Int32Array> & Record<K, Float32Array>` matching what the native side
actually produces (was incorrectly typed as `number[]`).
- **ARGMAX always returned**: Removed `'ARGMAX'` from
`classesOfInterest` — it's always in the output regardless, and the
return type reflects this.

### Introduces a breaking change?

- [x] Yes
- [ ] No

### Type of change

- [ ] Bug fix (change which fixes an issue)
- [x] New feature (change which adds functionality)
- [ ] Documentation update (improves or adds clarity to existing
documentation)
- [ ] Other (chores, tests, code style improvements etc.)

### Tested on

- [ ] iOS
- [x] Android

### Testing instructions

1. Build and run the `computer-vision` demo app
2. Navigate to Image Segmentation screen
3. Pick an image and run segmentation — verify the ARGMAX overlay
renders correctly
4. Verify the hook API works as expected:
```ts
const { isReady, forward } = useImageSegmentation({
  model: { modelName: 'deeplab-v3', modelSource: DEEPLAB_V3_RESNET50 },
});

// Returns Record<'ARGMAX', Int32Array> — no generic needed
const result = await forward(imageUri);

// Narrows return type to include 'PERSON' key as Float32Array
const result2 = await forward(imageUri, ['PERSON']);
```
5. Verify TypeScript autocompletion: `classesOfInterest` should only
suggest valid label keys for the chosen model (e.g. `'PERSON'`, `'CAR'`
for DeepLab, `'SELFIE'`/`'BACKGROUND'` for selfie segmentation)
6. You can also try changing the parameters, to say selfie segmentation
and see how the return types react. Please contact me for weights for
selfie segmentation as I'm not pushing them to HF yet

### Screenshots

<!-- Add screenshots here, if applicable -->

### Related issues

<!-- Link related issues here using #issue-number -->

### Checklist

- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have updated the documentation accordingly
- [x] My changes generate no new warnings

### Additional notes

The `ImageSegmentationModule.fromCustomConfig()` API allows users to
bring their own segmentation model with a custom label map:
```ts
const MyLabels = { BACKGROUND: 0, FOREGROUND: 1 } as const;
const seg = await ImageSegmentationModule.fromCustomConfig(
  'https://example.com/model.pte',
  { labelMap: MyLabels },
);
```

---------

Co-authored-by: Mateusz Kopcinski <120639731+mkopcins@users.noreply.github.com>

Author: chmjkb

.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.