add type for hook return

Author: chmjkb

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

@@ -1,7 +1,10 @@
 import Spinner from '../../components/Spinner';
 import { BottomBar } from '../../components/BottomBar';
 import { getImage } from '../../utils';
-import { useImageSegmentation } from 'react-native-executorch';
+import {
+  DEEPLAB_V3_RESNET50,
+  useImageSegmentation,
+} from 'react-native-executorch';
 import {
   Canvas,
   Image as SkiaImage,
@@ -43,10 +46,7 @@ export default function ImageSegmentationScreen() {
   const { setGlobalGenerating } = useContext(GeneratingContext);
   const { isReady, isGenerating, downloadProgress, forward } =
     useImageSegmentation({
-      model: {
-        modelName: 'deeplab-v3',
-        modelSource: 'https://ai.swmansion.com/storage/jc_tests/selfie_seg.pte',
-      },
+      model: DEEPLAB_V3_RESNET50,
     });
   const [imageUri, setImageUri] = useState('');
   const [imageSize, setImageSize] = useState({ width: 0, height: 0 });
@@ -72,7 +72,7 @@ export default function ImageSegmentationScreen() {
     if (!imageUri || imageSize.width === 0 || imageSize.height === 0) return;
     try {
       const { width, height } = imageSize;
-      const output = await forward(imageUri, ['PERSON'], true);
+      const output = await forward(imageUri, [], true);
       const argmax = output['ARGMAX'] || [];
       const pixels = new Uint8Array(width * height * 4);
 

packages/react-native-executorch/src/hooks/computer_vision/useImageSegmentation.ts

@@ -5,6 +5,7 @@ import {
 } from '../../modules/computer_vision/ImageSegmentationModule';
 import {
   ImageSegmentationProps,
+  ImageSegmentationType,
   ModelNameOf,
   ModelSources,
 } from '../../types/imageSegmentation';
@@ -30,7 +31,9 @@ import { RnExecutorchError, parseUnknownError } from '../../errors/errorUtils';
 export const useImageSegmentation = <C extends ModelSources>({
   model,
   preventLoad = false,
-}: ImageSegmentationProps<C>) => {
+}: ImageSegmentationProps<C>): ImageSegmentationType<
+  SegmentationLabels<ModelNameOf<C>>
+> => {
   const [error, setError] = useState<RnExecutorchError | null>(null);
   const [isReady, setIsReady] = useState(false);
   const [isGenerating, setIsGenerating] = useState(false);

packages/react-native-executorch/src/types/imageSegmentation.ts

@@ -1,3 +1,4 @@
+import { RnExecutorchError } from '../errors/errorUtils';
 import { LabelEnum, Triple, ResourceSource } from './common';
 
 /**
@@ -95,3 +96,47 @@ export interface ImageSegmentationProps<C extends ModelSources> {
   model: C;
   preventLoad?: boolean;
 }
+
+/**
+ * Return type for the `useImageSegmentation` hook.
+ * Manages the state and operations for image segmentation models.
+ *
+ * @typeParam L - The {@link LabelEnum} representing the model's class labels.
+ *
+ * @category Types
+ */
+export interface ImageSegmentationType<L extends LabelEnum> {
+  /**
+   * Contains the error object if the model failed to load, download, or encountered a runtime error during segmentation.
+   */
+  error: RnExecutorchError | null;
+
+  /**
+   * Indicates whether the segmentation model is loaded and ready to process images.
+   */
+  isReady: boolean;
+
+  /**
+   * Indicates whether the model is currently processing an image.
+   */
+  isGenerating: boolean;
+
+  /**
+   * Represents the download progress of the model binary as a value between 0 and 1.
+   */
+  downloadProgress: number;
+
+  /**
+   * Executes the model's forward pass to perform semantic segmentation on the provided image.
+   * @param imageSource - A string representing the image source (e.g., a file path, URI, or base64 string) to be processed.
+   * @param classesOfInterest - An optional array of label keys indicating which per-class probability masks to include in the output. `ARGMAX` is always returned regardless.
+   * @param resizeToInput - Whether to resize the output masks to the original input image dimensions. If `false`, returns the raw model output dimensions. Defaults to `true`.
+   * @returns A Promise resolving to an object with an `'ARGMAX'` `Int32Array` of per-pixel class indices, and each requested class label mapped to a `Float32Array` of per-pixel probabilities.
+   * @throws {RnExecutorchError} If the model is not loaded or is currently processing another image.
+   */
+  forward: <K extends keyof L>(
+    imageSource: string,
+    classesOfInterest?: K[],
+    resizeToInput?: boolean
+  ) => Promise<Record<'ARGMAX', Int32Array> & Record<K, Float32Array>>;
+}