Return class index, tweak mask drawing performance

Author: benITo47

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

@@ -1,10 +1,7 @@
 import Spinner from '../../components/Spinner';
 import { BottomBar } from '../../components/BottomBar';
 import { getImage } from '../../utils';
-import {
-  useInstanceSegmentation,
-  YOLO26X_SEG,
-} from 'react-native-executorch';
+import { useInstanceSegmentation, YOLO26N_SEG } from 'react-native-executorch';
 import {
   View,
   StyleSheet,
@@ -15,7 +12,10 @@ import {
 import React, { useContext, useEffect, useState } from 'react';
 import { GeneratingContext } from '../../context';
 import ScreenWrapper from '../../ScreenWrapper';
-import ImageWithMasks from '../../components/ImageWithMasks';
+import ImageWithMasks, {
+  buildDisplayInstances,
+  DisplayInstance,
+} from '../../components/ImageWithMasks';
 
 const AVAILABLE_INPUT_SIZES = [384, 512, 640];
 
@@ -24,12 +24,12 @@ export default function InstanceSegmentationScreen() {
 
   const { isReady, isGenerating, downloadProgress, forward, error } =
     useInstanceSegmentation({
-      model: YOLO26X_SEG,
+      model: YOLO26N_SEG,
     });
 
   const [imageUri, setImageUri] = useState('');
   const [imageSize, setImageSize] = useState({ width: 0, height: 0 });
-  const [instances, setInstances] = useState<any[]>([]);
+  const [instances, setInstances] = useState<DisplayInstance[]>([]);
   const [selectedInputSize, setSelectedInputSize] = useState(
     AVAILABLE_INPUT_SIZES[0]
   );
@@ -61,7 +61,10 @@ export default function InstanceSegmentationScreen() {
         inputSize: selectedInputSize,
       });
 
-      setInstances(output);
+      // Convert raw masks → small Skia images immediately.
+      // Raw Uint8Array mask buffers (backed by native OwningArrayBuffer)
+      // go out of scope here and become eligible for GC right away.
+      setInstances(buildDisplayInstances(output));
     } catch (e) {
       console.error(e);
     }

apps/computer-vision/components/ImageWithMasks.tsx

@@ -1,4 +1,4 @@
-import React, { useState } from 'react';
+import React, { useEffect, useState } from 'react';
 import { Image, StyleSheet, View, Text } from 'react-native';
 import {
   Canvas,
@@ -10,8 +10,6 @@ import {
   Rect,
   Group,
 } from '@shopify/react-native-skia';
-import type { SegmentedInstance } from 'react-native-executorch';
-import type { LabelEnum } from 'react-native-executorch';
 
 const INSTANCE_COLORS = [
   [255, 87, 51, 180],
@@ -26,39 +24,101 @@ const INSTANCE_COLORS = [
   [131, 51, 255, 180],
 ];
 
-interface Props {
-  imageUri: string;
-  instances: SegmentedInstance<LabelEnum>[];
-  imageWidth: number;
-  imageHeight: number;
+const MAX_MASK_DIM = 256;
+
+/** Display-only data — no raw mask buffers. */
+export interface DisplayInstance {
+  bbox: { x1: number; y1: number; x2: number; y2: number };
+  label: string;
+  score: number;
+  maskImage: SkImage;
+}
+
+/**
+ * Convert raw segmentation output into lightweight display instances.
+ * Call this eagerly (in the forward callback) so raw Uint8Array masks
+ * can be garbage-collected immediately.
+ */
+export function buildDisplayInstances(
+  rawInstances: {
+    bbox: { x1: number; y1: number; x2: number; y2: number };
+    mask: Uint8Array;
+    maskWidth: number;
+    maskHeight: number;
+    label: string | number;
+    score: number;
+  }[]
+): DisplayInstance[] {
+  return rawInstances
+    .map((inst, i) => {
+      const color = INSTANCE_COLORS[i % INSTANCE_COLORS.length];
+      const img = createMaskImage(
+        inst.mask,
+        inst.maskWidth,
+        inst.maskHeight,
+        color
+      );
+      if (!img) return null;
+      return {
+        bbox: inst.bbox,
+        label: String(inst.label),
+        score: inst.score,
+        maskImage: img,
+      };
+    })
+    .filter((d): d is DisplayInstance => d !== null);
 }
 
 function createMaskImage(
   mask: Uint8Array,
-  width: number,
-  height: number,
+  srcW: number,
+  srcH: number,
   color: number[]
 ): SkImage | null {
-  const pixels = new Uint8Array(width * height * 4);
-  for (let j = 0; j < mask.length; j++) {
-    if (mask[j] > 0) {
-      pixels[j * 4] = color[0];
-      pixels[j * 4 + 1] = color[1];
-      pixels[j * 4 + 2] = color[2];
-      pixels[j * 4 + 3] = color[3];
+  const downscale = Math.min(1, MAX_MASK_DIM / Math.max(srcW, srcH));
+  const dstW = Math.max(1, Math.round(srcW * downscale));
+  const dstH = Math.max(1, Math.round(srcH * downscale));
+
+  const pixels = new Uint8Array(dstW * dstH * 4);
+  const r = color[0],
+    g = color[1],
+    b = color[2],
+    a = color[3];
+
+  for (let dy = 0; dy < dstH; dy++) {
+    const sy = Math.min(Math.floor(dy / downscale), srcH - 1);
+    for (let dx = 0; dx < dstW; dx++) {
+      const sx = Math.min(Math.floor(dx / downscale), srcW - 1);
+      if (mask[sy * srcW + sx] > 0) {
+        const idx = (dy * dstW + dx) * 4;
+        pixels[idx] = r;
+        pixels[idx + 1] = g;
+        pixels[idx + 2] = b;
+        pixels[idx + 3] = a;
+      }
     }
   }
+
   const data = Skia.Data.fromBytes(pixels);
-  return Skia.Image.MakeImage(
+  const image = Skia.Image.MakeImage(
     {
-      width,
-      height,
+      width: dstW,
+      height: dstH,
       alphaType: AlphaType.Premul,
       colorType: ColorType.RGBA_8888,
     },
     data,
-    width * 4
+    dstW * 4
   );
+  data.dispose();
+  return image;
+}
+
+interface Props {
+  imageUri: string;
+  instances: DisplayInstance[];
+  imageWidth: number;
+  imageHeight: number;
 }
 
 export default function ImageWithMasks({
@@ -75,17 +135,12 @@ export default function ImageWithMasks({
   const offsetX = (layout.width - imageWidth * scale) / 2;
   const offsetY = (layout.height - imageHeight * scale) / 2;
 
-  const maskImages = instances
-    .map((instance, i) => {
-      const color = INSTANCE_COLORS[i % INSTANCE_COLORS.length];
-      return createMaskImage(
-        instance.mask,
-        instance.maskWidth,
-        instance.maskHeight,
-        color
-      );
-    })
-    .filter((img): img is SkImage => img !== null);
+  // Dispose Skia images when instances are replaced or on unmount
+  useEffect(() => {
+    return () => {
+      instances.forEach((inst) => inst.maskImage.dispose());
+    };
+  }, [instances]);
 
   return (
     <View
@@ -108,16 +163,15 @@ export default function ImageWithMasks({
       {instances.length > 0 && (
         <View style={styles.overlay}>
           <Canvas style={styles.canvas}>
-            {maskImages.map((maskImg, idx) => {
-              const inst = instances[idx];
+            {instances.map((inst, idx) => {
               const mx = inst.bbox.x1 * scale + offsetX;
               const my = inst.bbox.y1 * scale + offsetY;
               const mw = (inst.bbox.x2 - inst.bbox.x1) * scale;
               const mh = (inst.bbox.y2 - inst.bbox.y1) * scale;
               return (
                 <SkiaImage
                   key={`mask-${idx}`}
-                  image={maskImg}
+                  image={inst.maskImage}
                   fit="fill"
                   x={mx}
                   y={my}
@@ -127,12 +181,12 @@ export default function ImageWithMasks({
               );
             })}
 
-            {instances.map((instance, idx) => {
+            {instances.map((inst, idx) => {
               const color = INSTANCE_COLORS[idx % INSTANCE_COLORS.length];
-              const bboxX = instance.bbox.x1 * scale + offsetX;
-              const bboxY = instance.bbox.y1 * scale + offsetY;
-              const bboxW = (instance.bbox.x2 - instance.bbox.x1) * scale;
-              const bboxH = (instance.bbox.y2 - instance.bbox.y1) * scale;
+              const bboxX = inst.bbox.x1 * scale + offsetX;
+              const bboxY = inst.bbox.y1 * scale + offsetY;
+              const bboxW = (inst.bbox.x2 - inst.bbox.x1) * scale;
+              const bboxH = (inst.bbox.y2 - inst.bbox.y1) * scale;
 
               return (
                 <Group key={`bbox-${idx}`}>
@@ -150,10 +204,10 @@ export default function ImageWithMasks({
             })}
           </Canvas>
 
-          {instances.map((instance, idx) => {
+          {instances.map((inst, idx) => {
             const color = INSTANCE_COLORS[idx % INSTANCE_COLORS.length];
-            const bboxX = instance.bbox.x1 * scale + offsetX;
-            const bboxY = instance.bbox.y1 * scale + offsetY;
+            const bboxX = inst.bbox.x1 * scale + offsetX;
+            const bboxY = inst.bbox.y1 * scale + offsetY;
 
             return (
               <View
@@ -168,8 +222,7 @@ export default function ImageWithMasks({
                 ]}
               >
                 <Text style={styles.labelText}>
-                  {String(instance.label) || 'Unknown'}{' '}
-                  {(instance.score * 100).toFixed(0)}%
+                  {inst.label || 'Unknown'} {(inst.score * 100).toFixed(0)}%
                 </Text>
               </View>
             );

packages/react-native-executorch/common/rnexecutorch/host_objects/JsiConversions.h

@@ -480,9 +480,7 @@ inline jsi::Value getJsiValue(
     instance.setProperty(runtime, "maskWidth", instances[i].maskWidth);
     instance.setProperty(runtime, "maskHeight", instances[i].maskHeight);
 
-    instance.setProperty(
-        runtime, "label",
-        jsi::String::createFromUtf8(runtime, instances[i].label));
+    instance.setProperty(runtime, "classIndex", instances[i].classIndex);
 
     instance.setProperty(runtime, "score", instances[i].score);
     instance.setProperty(runtime, "instanceId", instances[i].instanceId);

packages/react-native-executorch/common/rnexecutorch/models/instance_segmentation/BaseInstanceSegmentation.cpp

@@ -10,10 +10,8 @@ namespace rnexecutorch::models::instance_segmentation {
 BaseInstanceSegmentation::BaseInstanceSegmentation(
     const std::string &modelSource, std::vector<float> normMean,
     std::vector<float> normStd, bool applyNMS,
-    std::vector<std::string> labelNames,
     std::shared_ptr<react::CallInvoker> callInvoker)
-    : BaseModel(modelSource, callInvoker), applyNMS_(applyNMS),
-      labelNames_(std::move(labelNames)) {
+    : BaseModel(modelSource, callInvoker), applyNMS_(applyNMS) {
   avalivableMethods_ = *module_->method_names();
   if (normMean.size() == 3) {
     normMean_ = cv::Scalar(normMean[0], normMean[1], normMean[2]);
@@ -66,7 +64,7 @@ std::vector<types::InstanceMask> BaseInstanceSegmentation::nonMaxSuppression(
         continue;
       }
 
-      if (instances[i].label == instances[j].label) {
+      if (instances[i].classIndex == instances[j].classIndex) {
         float iou = intersectionOverUnion(instances[i], instances[j]);
         if (iou > iouThreshold) {
           suppressed[j] = true;
@@ -150,14 +148,6 @@ std::vector<types::InstanceMask> BaseInstanceSegmentation::postprocess(
                                        labelIdx)) == allowedClasses.end())
       continue;
 
-    if (labelIdx >= labelNames_.size()) {
-      throw RnExecutorchError(
-          RnExecutorchErrorCode::InvalidConfig,
-          "Model output class index " + std::to_string(labelIdx) +
-              " exceeds labelNames size " + std::to_string(labelNames_.size()) +
-              ". Ensure the labelMap covers all model output classes.");
-    }
-
     // Scale bbox to original image coordinates
     float origX1 = x1 * widthRatio;
     float origY1 = y1 * heightRatio;
@@ -246,15 +236,15 @@ std::vector<types::InstanceMask> BaseInstanceSegmentation::postprocess(
     instance.mask = std::move(finalMask);
     instance.maskWidth = finalMaskWidth;
     instance.maskHeight = finalMaskHeight;
-    instance.label = labelNames_[labelIdx];
+    instance.classIndex = static_cast<int32_t>(labelIdx);
     instance.score = score;
     instance.instanceId = i;
     instances.push_back(std::move(instance));
     ++processed;
   }
 
   // Finalize: NMS + limit + renumber
-  if (applyNMS_ && iouThreshold < 0.45) {
+  if (applyNMS_) {
     instances = nonMaxSuppression(instances, iouThreshold);
   }
 

packages/react-native-executorch/common/rnexecutorch/models/instance_segmentation/BaseInstanceSegmentation.h

@@ -20,7 +20,6 @@ class BaseInstanceSegmentation : public BaseModel {
   BaseInstanceSegmentation(const std::string &modelSource,
                            std::vector<float> normMean,
                            std::vector<float> normStd, bool applyNMS,
-                           std::vector<std::string> labelNames,
                            std::shared_ptr<react::CallInvoker> callInvoker);
 
   [[nodiscard("Registered non-void function")]] std::vector<types::InstanceMask>
@@ -48,7 +47,6 @@ class BaseInstanceSegmentation : public BaseModel {
   std::optional<cv::Scalar> normMean_;
   std::optional<cv::Scalar> normStd_;
   bool applyNMS_;
-  std::vector<std::string> labelNames_;
   cv::Size modelImageSize{0, 0};
   std::unordered_set<std::string> avalivableMethods_;
   std::string currentlyLoadedMethod_;
@@ -57,6 +55,5 @@ class BaseInstanceSegmentation : public BaseModel {
 
 REGISTER_CONSTRUCTOR(models::instance_segmentation::BaseInstanceSegmentation,
                      std::string, std::vector<float>, std::vector<float>, bool,
-                     std::vector<std::string>,
                      std::shared_ptr<react::CallInvoker>);
 } // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/models/instance_segmentation/Types.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <string>
+#include <cstdint>
 #include <vector>
 
 namespace rnexecutorch::models::instance_segmentation::types {
@@ -19,7 +19,7 @@ struct InstanceMask {
   std::vector<uint8_t> mask; ///< Binary mask (0 or 1) for the instance
   int maskWidth;             ///< Width of the mask array
   int maskHeight;            ///< Height of the mask array
-  std::string label;         ///< Class label name
+  int32_t classIndex;        ///< Model output class index
   float score;               ///< Confidence score [0, 1]
   int instanceId;            ///< Unique identifier for this instance
 };