[RNE Rewrite] feat: add semantic segmentation task

apps/computer-vision/app/_layout.tsx

@@ -27,6 +27,13 @@ export default function Layout() {
           title: 'Classification',
         }}
       />
+      <Drawer.Screen
+        name="segmentation/index"
+        options={{
+          drawerLabel: 'Semantic Segmentation',
+          title: 'Semantic Segmentation',
+        }}
+      />
       <Drawer.Screen
         name="inspect/index"
         options={{

apps/computer-vision/app/index.tsx

@@ -14,6 +14,9 @@ export default function Home() {
         <TouchableOpacity style={styles.button} onPress={() => router.navigate('classification/')}>
           <Text style={styles.buttonText}>Classification</Text>
         </TouchableOpacity>
+        <TouchableOpacity style={styles.button} onPress={() => router.navigate('segmentation/')}>
+          <Text style={styles.buttonText}>Semantic Segmentation</Text>
+        </TouchableOpacity>
         <TouchableOpacity style={styles.button} onPress={() => router.navigate('inspect/')}>
           <Text style={styles.buttonText}>Model Inspector</Text>
         </TouchableOpacity>

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

@@ -0,0 +1,206 @@
+import React, { useState, useRef } from 'react';
+import { View, Text, ScrollView } from 'react-native';
+import { useSafeAreaInsets } from 'react-native-safe-area-context';
+import { commonStyles, theme } from '../../theme';
+import {
+  Skia,
+  ColorType,
+  AlphaType,
+  useImage,
+  type SkImage as SkiaImageType,
+} from '@shopify/react-native-skia';
+import { useSemanticSegmenter, models } from 'react-native-executorch';
+import ScreenWrapper from '../../components/ScreenWrapper';
+import { ModelPicker, type ModelOption } from '../../components/ModelPicker';
+import { getImage } from '../../utils';
+import { ImageViewport } from '../../components/ImageViewport';
+import { ModelStatus } from '../../components/ModelStatus';
+import { LatencyIndicator } from '../../components/LatencyIndicator';
+import { Button } from '../../components/Button';
+
+const SEGMENTATION_OPTIONS: ModelOption[] = [
+  {
+    label: 'Selfie Seg (FP32)',
+    value: models.semanticSegmentation.SELFIE_SEGMENTATION.XNNPACK_FP32,
+  },
+  {
+    label: 'LRASPP MobileNet V3 (INT8)',
+    value: models.semanticSegmentation.LRASPP_MOBILENET_V3_LARGE.XNNPACK_INT8,
+  },
+  {
+    label: 'DeepLab V3 ResNet50 (INT8)',
+    value: models.semanticSegmentation.DEEPLAB_V3_RESNET50.XNNPACK_INT8,
+  },
+  {
+    label: 'DeepLab V3 ResNet101 (INT8)',
+    value: models.semanticSegmentation.DEEPLAB_V3_RESNET101.XNNPACK_INT8,
+  },
+  {
+    label: 'DeepLab V3 MobileNet V3 (INT8)',
+    value: models.semanticSegmentation.DEEPLAB_V3_MOBILENET_V3_LARGE.XNNPACK_INT8,
+  },
+  {
+    label: 'FCN ResNet50 (INT8)',
+    value: models.semanticSegmentation.FCN_RESNET50.XNNPACK_INT8,
+  },
+  {
+    label: 'FCN ResNet101 (INT8)',
+    value: models.semanticSegmentation.FCN_RESNET101.XNNPACK_INT8,
+  },
+];
+
+function SegmentationContent() {
+  const insets = useSafeAreaInsets();
+  const [selectedModel, setSelectedModel] = useState<any>(SEGMENTATION_OPTIONS[0].value);
+  const [imageUri, setImageUri] = useState<string | null>(null);
+  const [isProcessing, setIsProcessing] = useState(false);
+  const [latency, setLatency] = useState<number | null>(null);
+  const [segmentationImage, setSegmentationImage] = useState<SkiaImageType | null>(null);
+  const [error, setError] = useState<string | null>(null);
+
+  const isProcessingRef = useRef(false);
+
+  const skiaImage = useImage(imageUri, (err) => setError(err.message || String(err)));
+
+  const {
+    isReady,
+    downloadProgress,
+    error: loadError,
+    segment,
+    segmentWorklet,
+  } = useSemanticSegmenter(selectedModel);
+
+  const handlePickImage = async (useCamera: boolean) => {
+    setError(null);
+    try {
+      const uri = await getImage(useCamera);
+      if (uri) {
+        setImageUri(uri);
+        setLatency(null);
+        setSegmentationImage(null);
+      }
+    } catch (e: any) {
+      setError(e.message || String(e));
+    }
+  };
+
+  const runSegmentation = async (sync: boolean) => {
+    if (isProcessingRef.current) return;
+    if (!skiaImage || !segment || !segmentWorklet) return;
+
+    isProcessingRef.current = true;
+    setIsProcessing(true);
+    setError(null);
+
+    try {
+      const pixels = skiaImage.readPixels();
+      if (!pixels) {
+        throw new Error('Failed to read pixels from image');
+      }
+      if (!(pixels instanceof Uint8Array)) {
+        throw new Error('Expected Uint8Array from readPixels');
+      }
+      const buffer = {
+        data: pixels,
+        width: skiaImage.width(),
+        height: skiaImage.height(),
+        format: 'rgba' as const,
+        layout: 'hwc' as const,
+      };
+
+      const start = Date.now();
+      const { buffer: outBuffer } = sync ? segmentWorklet(buffer) : await segment(buffer);
+      setLatency(Date.now() - start);
+
+      const outData = Skia.Data.fromBytes(outBuffer.data);
+      const info = {
+        width: buffer.width,
+        height: buffer.height,
+        colorType: ColorType.RGBA_8888,
+        alphaType: AlphaType.Premul,
+      };
+      const nextImage = Skia.Image.MakeImage(info, outData, buffer.width * 4);
+      if (!nextImage) {
+        throw new Error('Failed to create overlay image from output data');
+      }
+      setSegmentationImage(nextImage);
+    } catch (e: any) {
+      setError(e.message || String(e));
+    } finally {
+      isProcessingRef.current = false;
+      setIsProcessing(false);
+    }
+  };
+
+  const activeError = loadError ? String(loadError) : error;
+
+  return (
+    <ScrollView
+      style={commonStyles.container}
+      contentContainerStyle={[
+        commonStyles.contentContainer,
+        { paddingBottom: insets.bottom + theme.spacing.large },
+      ]}
+    >
+      <Text style={commonStyles.description}>
+        Upload or capture an image to partition it into multiple segments using semantic
+        segmentation.
+      </Text>
+
+      <ModelPicker
+        label="Model"
+        options={SEGMENTATION_OPTIONS}
+        selectedValue={selectedModel}
+        onValueChange={(model) => {
+          setSelectedModel(model);
+          setLatency(null);
+          setError(null);
+          setSegmentationImage(null);
+        }}
+      />
+
+      <ModelStatus
+        isReady={isReady}
+        downloadProgress={downloadProgress}
+        error={activeError}
+        modelTypeLabel="segmentation model"
+      />
+
+      <ImageViewport
+        skiaImage={skiaImage}
+        overlayImage={segmentationImage}
+        onPressPlaceholder={() => handlePickImage(false)}
+      />
+
+      <View style={commonStyles.buttonRow}>
+        <Button title="Gallery" onPress={() => handlePickImage(false)} variant="secondary" />
+        <Button title="Camera" onPress={() => handlePickImage(true)} variant="secondary" />
+      </View>
+
+      <View style={commonStyles.buttonRow}>
+        <Button
+          title="Run Async"
+          onPress={() => runSegmentation(false)}
+          disabled={!skiaImage || !isReady || isProcessing}
+          loading={isProcessing}
+        />
+        <Button
+          title="Run Sync"
+          onPress={() => runSegmentation(true)}
+          disabled={!skiaImage || !isReady || isProcessing}
+          variant="accent"
+        />
+      </View>
+
+      <LatencyIndicator latency={latency} />
+    </ScrollView>
+  );
+}
+
+export default function SegmentationScreen() {
+  return (
+    <ScreenWrapper>
+      <SegmentationContent />
+    </ScreenWrapper>
+  );
+}

packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp

@@ -614,4 +614,93 @@ void install_normalize(jsi::Runtime &rt, jsi::Object &module) {
 
     module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 3, fnBody));
 }
+
+void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) {
+    auto name = "applyColormap";
+    auto fnBody = [](jsi::Runtime &rt, const jsi::Value &thisVal, const jsi::Value *args, size_t count) -> jsi::Value {
+        if (count != 3) {
+            throw jsi::JSError(rt, "Usage: applyColormap(src, dst, colormap)");
+        }
+
+        if (!args[0].isObject() || !args[0].asObject(rt).isHostObject<TensorHostObject>(rt)) {
+            throw jsi::JSError(rt, "applyColormap: src must be a Tensor");
+        }
+        if (!args[1].isObject() || !args[1].asObject(rt).isHostObject<TensorHostObject>(rt)) {
+            throw jsi::JSError(rt, "applyColormap: dst must be a Tensor");
+        }
+
+        auto src = args[0].asObject(rt).getHostObject<TensorHostObject>(rt);
+        auto dst = args[1].asObject(rt).getHostObject<TensorHostObject>(rt);
+        constexpr size_t numRgbaChannels = 4;
+
+        if (src->dtype_ != rnexecutorch::core::types::DType::int32) {
+            throw jsi::JSError(rt, "applyColormap: src must be int32");
+        }
+        if (dst->dtype_ != rnexecutorch::core::types::DType::uint8) {
+            throw jsi::JSError(rt, "applyColormap: dst must be uint8");
+        }
+        if (dst->numel_ != src->numel_ * numRgbaChannels) {
+            throw jsi::JSError(rt, "applyColormap: dst must have exactly 4 times the number of elements as src (RGBA channels)");
+        }
+
+        if (!args[2].isObject() || !args[2].asObject(rt).isArray(rt)) {
+            throw jsi::JSError(rt, "applyColormap: colormap must be an array");
+        }
+
+        auto colormapArray = args[2].asObject(rt).asArray(rt);
+        size_t numColors = colormapArray.size(rt);
+        std::vector<std::array<uint8_t, numRgbaChannels>> lut(numColors);
+        for (size_t i = 0; i < numColors; ++i) {
+            auto colorVal = colormapArray.getValueAtIndex(rt, i);
+            if (!colorVal.isObject() || !colorVal.asObject(rt).isArray(rt)) {
+                throw jsi::JSError(rt, "applyColormap: colormap entry must be an array");
+            }
+            auto color = colorVal.asObject(rt).asArray(rt);
+            if (color.size(rt) != numRgbaChannels) {
+                throw jsi::JSError(rt, "applyColormap: colormap entry must be an RGBA color array of size 4");
+            }
+            for (size_t c = 0; c < numRgbaChannels; ++c) {
+                auto channelVal = color.getValueAtIndex(rt, c);
+                if (!channelVal.isNumber()) {
+                    throw jsi::JSError(rt, "applyColormap: colormap channel value must be a number");
+                }
+                double val = channelVal.asNumber();
+                if (std::isnan(val) || val < 0.0 || val > 255.0) {
+                    throw jsi::JSError(rt, "applyColormap: colormap channel value must be between 0 and 255");
+                }
+                lut[i][c] = static_cast<uint8_t>(val);
+            }
+        }
+
+        std::shared_lock<std::shared_mutex> srcLock(src->mutex_, std::try_to_lock);
+        std::unique_lock<std::shared_mutex> dstLock(dst->mutex_, std::try_to_lock);
+        if (!srcLock.owns_lock() || !dstLock.owns_lock()) {
+            throw jsi::JSError(rt, "applyColormap: tensors in use");
+        }
+
+        if (!src->data_ || !dst->data_) {
+            throw jsi::JSError(rt, "applyColormap: tensor has been disposed");
+        }
+
+        size_t pixels = src->numel_;
+
+        const int32_t *srcData = reinterpret_cast<const int32_t *>(src->data_.get());
+        uint8_t *dstData = dst->data_.get();
+
+        for (size_t i = 0; i < pixels; ++i) {
+            int32_t idx = srcData[i];
+            if (idx < 0 || static_cast<size_t>(idx) >= numColors) {
+                throw jsi::JSError(rt, "applyColormap: tensor contains class index (" +
+                                           std::to_string(idx) + ") that exceeds provided colormap size (" +
+                                           std::to_string(numColors) + ")");
+            }
+            for (size_t c = 0; c < numRgbaChannels; ++c) {
+                dstData[i * numRgbaChannels + c] = lut[idx][c];
+            }
+        }
+
+        return jsi::Value(rt, args[1]);
+    };
+    module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 3, fnBody));
+}
 } // namespace rnexecutorch::extensions::cv::image_ops

packages/react-native-executorch/cpp/extensions/cv/image_ops.h

@@ -8,4 +8,5 @@ void install_cvtColor(facebook::jsi::Runtime &rt, facebook::jsi::Object &module)
 void install_toChannelsFirst(facebook::jsi::Runtime &rt, facebook::jsi::Object &module);
 void install_toChannelsLast(facebook::jsi::Runtime &rt, facebook::jsi::Object &module);
 void install_normalize(facebook::jsi::Runtime &rt, facebook::jsi::Object &module);
+void install_applyColormap(facebook::jsi::Runtime &rt, facebook::jsi::Object &module);
 } // namespace rnexecutorch::extensions::cv::image_ops

packages/react-native-executorch/cpp/extensions/cv/install.cpp

@@ -12,6 +12,7 @@ void install(facebook::jsi::Runtime &rt, facebook::jsi::Object &module) {
     image_ops::install_toChannelsFirst(rt, cvModule);
     image_ops::install_toChannelsLast(rt, cvModule);
     image_ops::install_normalize(rt, cvModule);
+    image_ops::install_applyColormap(rt, cvModule);
 
     module.setProperty(rt, "cv", cvModule);
 }

packages/react-native-executorch/cpp/extensions/math/operations.cpp

@@ -277,26 +277,21 @@ void install_argmax(jsi::Runtime &rt, jsi::Object &module) {
         }
 
         int32_t *dstData = reinterpret_cast<int32_t *>(dst->data_.get());
-        std::vector<float> maxVals(inner);
 
+        // DO NOT swap loop order. This structure intentionally prioritizes the
+        // most common case (axis = -1, inner = 1) for sequential access.
         for (size_t o = 0; o < outer; ++o) {
-            const float *srcSlab = srcData + o * axisDim * inner;
-            int32_t *dstRow = dstData + o * inner;
-
             for (size_t i = 0; i < inner; ++i) {
-                maxVals[i] = -std::numeric_limits<float>::infinity();
-                dstRow[i] = 0;
-            }
-
-            for (size_t d = 0; d < axisDim; ++d) {
-                const float *srcRow = srcSlab + d * inner;
-                for (size_t i = 0; i < inner; ++i) {
-                    const float val = srcRow[i];
-                    if (val > maxVals[i]) {
-                        maxVals[i] = val;
-                        dstRow[i] = static_cast<int32_t>(d);
+                float maxVal = -std::numeric_limits<float>::infinity();
+                int32_t maxIdx = 0;
+                for (size_t d = 0; d < axisDim; ++d) {
+                    const float val = srcData[o * axisDim * inner + d * inner + i];
+                    if (val > maxVal) {
+                        maxVal = val;
+                        maxIdx = static_cast<int32_t>(d);
                     }
                 }
+                dstData[o * inner + i] = maxIdx;
             }
         }