From f0849020a461263c8a06ed82c749ca5df1e4247c Mon Sep 17 00:00:00 2001 From: Bartosz Hanc Date: Mon, 22 Jun 2026 16:28:55 +0200 Subject: [PATCH 1/3] feat: add semantic segmentation task and native image operations --- .../cpp/extensions/cv/image_ops.cpp | 65 ++++++ .../cpp/extensions/cv/image_ops.h | 1 + .../cpp/extensions/cv/install.cpp | 1 + .../cpp/extensions/math/operations.cpp | 23 +-- .../react-native-executorch/src/constants.ts | 34 +++ .../src/extensions/cv/ops/image.ts | 19 ++ .../cv/tasks/semanticSegmentation.ts | 195 ++++++++++++++++++ .../src/hooks/useSemanticSegmenter.ts | 47 +++++ packages/react-native-executorch/src/index.ts | 2 + .../react-native-executorch/src/models.ts | 51 ++++- 10 files changed, 422 insertions(+), 16 deletions(-) create mode 100644 packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts create mode 100644 packages/react-native-executorch/src/hooks/useSemanticSegmenter.ts diff --git a/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp b/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp index 7f5e4b06de..61ffc01704 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp +++ b/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp @@ -614,4 +614,69 @@ 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)"); + } + + auto src = args[0].asObject(rt).getHostObject(rt); + auto dst = args[1].asObject(rt).getHostObject(rt); + + 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"); + } + + auto colormapArray = args[2].asObject(rt).asArray(rt); + size_t numColors = colormapArray.size(rt); + std::vector> lut(numColors); + for (size_t i = 0; i < numColors; ++i) { + auto color = colormapArray.getValueAtIndex(rt, i).asObject(rt).asArray(rt); + lut[i][0] = color.getValueAtIndex(rt, 0).asNumber(); + lut[i][1] = color.getValueAtIndex(rt, 1).asNumber(); + lut[i][2] = color.getValueAtIndex(rt, 2).asNumber(); + lut[i][3] = color.getValueAtIndex(rt, 3).asNumber(); + } + + std::shared_lock src_lock(src->mutex_, std::try_to_lock); + std::unique_lock dst_lock(dst->mutex_, std::try_to_lock); + if (!src_lock.owns_lock() || !dst_lock.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 = 1; + for (int i = 0; i < src->shape_.size(); ++i) { + pixels *= src->shape_[i]; + } + + const int32_t *srcData = reinterpret_cast(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(idx) >= numColors) { + throw jsi::JSError(rt, "applyColormap: tensor contains class index (" + + std::to_string(idx) + ") that exceeds provided colormap size (" + + std::to_string(numColors) + ")"); + } + + dstData[i * 4 + 0] = lut[idx][0]; + dstData[i * 4 + 1] = lut[idx][1]; + dstData[i * 4 + 2] = lut[idx][2]; + dstData[i * 4 + 3] = lut[idx][3]; + } + + 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 diff --git a/packages/react-native-executorch/cpp/extensions/cv/image_ops.h b/packages/react-native-executorch/cpp/extensions/cv/image_ops.h index fe6cfa546d..893c0957e5 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/image_ops.h +++ b/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 diff --git a/packages/react-native-executorch/cpp/extensions/cv/install.cpp b/packages/react-native-executorch/cpp/extensions/cv/install.cpp index 676fab3cc0..f559c68552 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/install.cpp +++ b/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); } diff --git a/packages/react-native-executorch/cpp/extensions/math/operations.cpp b/packages/react-native-executorch/cpp/extensions/math/operations.cpp index b4c77a4adc..614404a73e 100644 --- a/packages/react-native-executorch/cpp/extensions/math/operations.cpp +++ b/packages/react-native-executorch/cpp/extensions/math/operations.cpp @@ -277,26 +277,19 @@ void install_argmax(jsi::Runtime &rt, jsi::Object &module) { } int32_t *dstData = reinterpret_cast(dst->data_.get()); - std::vector maxVals(inner); 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::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(d); + float maxVal = -std::numeric_limits::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(d); } } + dstData[o * inner + i] = maxIdx; } } diff --git a/packages/react-native-executorch/src/constants.ts b/packages/react-native-executorch/src/constants.ts index 789f7358e2..224bb09e3d 100644 --- a/packages/react-native-executorch/src/constants.ts +++ b/packages/react-native-executorch/src/constants.ts @@ -1005,8 +1005,42 @@ export const IMAGENET1K_LABELS = [ 'toilet tissue, toilet paper, bathroom tissue', ] as const; +/** + * Pascal VOC dataset label array containing the 21 categories. + * @category Constants + */ +export const PASCAL_VOC_LABELS = [ + 'background', + 'aeroplane', + 'bicycle', + 'bird', + 'boat', + 'bottle', + 'bus', + 'car', + 'cat', + 'chair', + 'cow', + 'diningtable', + 'dog', + 'horse', + 'motorbike', + 'person', + 'pottedplant', + 'sheep', + 'sofa', + 'train', + 'tvmonitor', +] as const; + /** * Type representing a valid ImageNet 1K label string. * @category Types */ export type ImageNet1KLabel = (typeof IMAGENET1K_LABELS)[number]; + +/** + * Type representing a valid Pascal VOC label string. + * @category Types + */ +export type PascalVocLabel = (typeof PASCAL_VOC_LABELS)[number]; diff --git a/packages/react-native-executorch/src/extensions/cv/ops/image.ts b/packages/react-native-executorch/src/extensions/cv/ops/image.ts index 352ad7f484..a1fd9671d6 100644 --- a/packages/react-native-executorch/src/extensions/cv/ops/image.ts +++ b/packages/react-native-executorch/src/extensions/cv/ops/image.ts @@ -175,3 +175,22 @@ export function normalize(src: Tensor, dst: Tensor, opts?: NormalizeOptions): Te } as const; return rnexecutorchJsi.cv.normalize(src, dst, { ...defaultNormalizeOptions, ...opts }); } + +/** + * Applies a colormap to a 2D single-channel image tensor, mapping class indices + * to RGBA colors. + * @category Typescript API + * @param src The source index/mask tensor (int32). + * @param dst The pre-allocated destination tensor (uint8, 4 channels). + * @param colormap An array of RGBA color arrays corresponding to the class + * indices. + * @returns The destination tensor with applied colormap. + */ +export function applyColormap( + src: Tensor, + dst: Tensor, + colormap: [number, number, number, number][] +): Tensor { + 'worklet'; + return rnexecutorchJsi.cv.applyColormap(src, dst, colormap); +} diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts new file mode 100644 index 0000000000..2bd1f4e4d7 --- /dev/null +++ b/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts @@ -0,0 +1,195 @@ +import type { WorkletRuntime } from 'react-native-worklets'; + +import { tensor } from '../../../core/tensor'; +import { loadModel } from '../../../core/model'; +import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema'; +import { wrapAsync } from '../../../core/runtime'; + +import { type ImageBuffer } from '../image'; +import { createImagePreprocessor, type ImagePreprocessorOptions } from './preprocessing'; +import { + toChannelsLast, + normalize, + resize, + cvtColor, + applyColormap, + type InterpolationMethod, +} from '../ops/image'; +import { sigmoid, argmax } from '../../math'; + +/** + * Options for configuring a semantic segmenter preprocessor and label + * vocabulary. + * @category Types + */ +export type SemanticSegmentationOptions = Omit & { + readonly resizeMode: 'stretch'; + readonly outInterpolation: InterpolationMethod; + readonly labels: readonly L[]; +}; + +/** + * Model configuration required to instantiate a segmenter task runner. + * @category Types + */ +export type SemanticSegmentationModel = { + readonly modelPath: string; + readonly opts: SemanticSegmentationOptions; +}; + +/** + * Maps each class label to its assigned RGBA color. + * @category Types + */ +export type ColorMap = Record; + +/** + * Result structure representing the output of a semantic segmentation task. + * @category Types + */ +export type SemanticSegmentationResult = { + buffer: ImageBuffer; + colormap?: ColorMap; +}; + +function hslToRgb(h: number, s: number, l: number): [number, number, number] { + s /= 100; + l /= 100; + const k = (n: number) => (n + h / 30) % 12; + const a = s * Math.min(l, 1 - l); + const f = (n: number) => l - a * Math.max(-1, Math.min(k(n) - 3, 9 - k(n), 1)); + return [Math.round(255 * f(0)), Math.round(255 * f(8)), Math.round(255 * f(4))]; +} + +/** + * Creates a semantic segmenter runner for executing local Semantic Segmentation + * models. + * + * It validates the model inputs and outputs, asserts that the labels array + * length matches the model's output vocabulary size, pre-allocates the + * necessary static execution tensors, sets up an image preprocessor, and + * registers clean disposal hooks to clear all native memory. + * @category Typescript API + * @typeParam L The type representing the segmentation labels. + * @param config Segmenter task configuration containing path and options. + * @param runtime Optional worklet runtime thread environment context. + * @returns A promise resolving to an object containing segmentation and + * disposal controls. + */ +export async function createSemanticSegmenter( + config: SemanticSegmentationModel, + runtime?: WorkletRuntime +): Promise<{ + dispose: () => void; + segment: ( + input: ImageBuffer, + colormap?: Partial> + ) => Promise>; + segmentWorklet: ( + input: ImageBuffer, + colormap?: Partial> + ) => SemanticSegmentationResult; +}> { + const { modelPath, opts } = config; + const model = await wrapAsync(loadModel, runtime)(modelPath); + + const meta = validateModelSchema( + model, + 'forward', + [SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])], + [SymbolicTensor('float32', [1, 'K', 'H', 'W'], ['K', 'H', 'W'])] + ); + const inpShape = meta.inputTensorMeta[0]!.shape; + const outShape = meta.outputTensorMeta[0]!.shape; + + const nClasses = outShape.at(-3)!; + const targetH = outShape.at(-2)!; + const targetW = outShape.at(-1)!; + + // Generate highly distinct, high-contrast colors, see: + // https://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically/ + const defaultColormap = opts.labels.map((_, i) => { + if (i === 0) return [0, 0, 0, 0] as const; + return [...hslToRgb((i * 137.5) % 360, 95, 50), 255] as const; + }); + + if (nClasses > 1 && opts.labels.length !== nClasses) { + throw new Error( + `Model outputs ${nClasses} classes, but ${opts.labels.length} labels were provided in the configuration.` + ); + } + + const tensors = [ + tensor('float32', outShape), + tensor('float32', [nClasses, targetH, targetW]), + tensor('float32', [nClasses, targetH, targetW]), + tensor('float32', [targetH, targetW, nClasses]), + tensor(nClasses > 1 ? 'int32' : 'uint8', [targetH, targetW, 1]), + tensor('uint8', [targetH, targetW, 4]), + ] as const; + + const [tOutput, tReshape, tSigmoid, tChanLast, tMask, tRgba] = tensors; + const preprocessor = createImagePreprocessor(opts, inpShape); + + const dispose = () => { + tensors.forEach((t) => t.dispose()); + preprocessor.dispose(); + model.dispose(); + }; + + const segmentWorklet = ( + input: ImageBuffer, + colormap?: Partial> + ): SemanticSegmentationResult => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tOutput]); + + let returnColormap: ColorMap | undefined; + if (nClasses > 1) { + if (colormap) { + returnColormap = Object.fromEntries( + opts.labels.map((l) => [l, colormap[l] ?? [0, 0, 0, 0]]) + ) as ColorMap; + } else { + returnColormap = Object.fromEntries( + opts.labels.map((l, i) => [l, defaultColormap[i]!]) + ) as ColorMap; + } + + const colormapData = opts.labels.map((l) => returnColormap![l]); + + tOutput + .copyTo(tReshape) + .through(toChannelsLast, tChanLast) + .through(argmax, tMask, -1) + .through(applyColormap, tRgba, colormapData); + } else { + tOutput + .copyTo(tReshape) + .through(sigmoid, tSigmoid) + .through(toChannelsLast, tChanLast) + .through(normalize, tMask, { alpha: 255.0 }) + .through(cvtColor, tRgba, 'GRAY2RGBA'); + } + + const data = new Uint8Array(input.width * input.height * 4); + const tResize = tensor('uint8', [input.height, input.width, 4]); + try { + tRgba + .through(resize, tResize, { mode: 'stretch', interpolation: opts.outInterpolation }) + .getData(data); + } finally { + tResize.dispose(); + } + + return { + buffer: { data, width: input.width, height: input.height, format: 'rgba', layout: 'hwc' }, + colormap: returnColormap, + }; + }; + + const segment = wrapAsync(segmentWorklet, runtime); + + return { segment, segmentWorklet, dispose }; +} diff --git a/packages/react-native-executorch/src/hooks/useSemanticSegmenter.ts b/packages/react-native-executorch/src/hooks/useSemanticSegmenter.ts new file mode 100644 index 0000000000..8140204e18 --- /dev/null +++ b/packages/react-native-executorch/src/hooks/useSemanticSegmenter.ts @@ -0,0 +1,47 @@ +import { useModel } from './useModel'; +import { useResourceDownload } from './useResourceDownload'; +import { + createSemanticSegmenter, + type SemanticSegmentationModel, +} from '../extensions/cv/tasks/semanticSegmentation'; + +/** + * React hook to load and run a semantic segmentation model. + * + * This hook manages downloading (if it's a remote URL) and loading the model + * file, compiling it, tracking download progress and compilation errors, and + * cleaning up native model memory when the component unmounts or configuration + * changes. + * @category Hooks + * @typeParam L The type representing the segmentation labels. + * @param config The semantic segmentation model configuration. + * @param options Hook options. + * @param options.preventLoad If true, prevents downloading and compiling the + * model. + * @returns An object containing the model's loading state, error, download + * progress, and segmentation functions. + */ +export function useSemanticSegmenter( + config: SemanticSegmentationModel, + options?: { preventLoad?: boolean } +) { + const { localPath, downloadProgress, downloadError } = useResourceDownload( + config.modelPath, + options?.preventLoad + ); + const { model, error } = useModel( + createSemanticSegmenter, + localPath ? { ...config, modelPath: localPath } : null, + [localPath] + ); + + return { + isReady: !!model, + error: downloadError || error, + downloadProgress, + localPath, + segment: model?.segment, + segmentWorklet: model?.segmentWorklet, + labels: config.opts.labels, + }; +} diff --git a/packages/react-native-executorch/src/index.ts b/packages/react-native-executorch/src/index.ts index 8f709b8433..760ca5467c 100644 --- a/packages/react-native-executorch/src/index.ts +++ b/packages/react-native-executorch/src/index.ts @@ -1,5 +1,6 @@ // Hooks — primary API for app developers export * from './hooks/useClassifier'; +export * from './hooks/useSemanticSegmenter'; export * from './hooks/useResourceDownload'; export * from './hooks/useModel'; @@ -9,6 +10,7 @@ export * as constants from './constants'; // Task APIs — for developers needing manual lifetime/disposal control export * from './extensions/cv/tasks/classification'; +export * from './extensions/cv/tasks/semanticSegmentation'; // Core primitives — for library builders and power users export { tensor } from './core/tensor'; diff --git a/packages/react-native-executorch/src/models.ts b/packages/react-native-executorch/src/models.ts index 66a88e937b..befe771647 100644 --- a/packages/react-native-executorch/src/models.ts +++ b/packages/react-native-executorch/src/models.ts @@ -1,5 +1,11 @@ import type { ClassifierModel } from './extensions/cv/tasks/classification'; -import { IMAGENET1K_LABELS, type ImageNet1KLabel } from './constants'; +import type { SemanticSegmentationModel } from './extensions/cv/tasks/semanticSegmentation'; +import { + IMAGENET1K_LABELS, + PASCAL_VOC_LABELS, + type ImageNet1KLabel, + type PascalVocLabel, +} from './constants'; const BASE_URL = 'https://huggingface.co/software-mansion/react-native-executorch'; const VERSION_TAG = 'resolve/v0.9.0'; @@ -27,6 +33,38 @@ const EFFICIENTNET_V2_S_COREML_FP16: ClassifierModel = { classifierOpts: EFFICIENTNET_V2_S_OPTS, }; +// ============================================================================= +// Semantic Segmentation +// ============================================================================= +const SELFIE_SEGMENTATION_XNNPACK_FP32: SemanticSegmentationModel<'background' | 'person'> = { + modelPath: `${BASE_URL}-selfie-segmentation/${VERSION_TAG}/xnnpack/selfie_segmentation_xnnpack_fp32.pte`, + opts: { + labels: ['background', 'person'] as const, + resizeMode: 'stretch', + interpolation: 'linear', + alpha: 1 / 255.0, + beta: 0.0, + outInterpolation: 'lanczos', + }, +}; + +const LRASPP_MOBILENET_V3_LARGE_OPTS = { + labels: PASCAL_VOC_LABELS, + resizeMode: 'stretch' as const, + interpolation: 'linear' as const, + alpha: [1 / (255.0 * 0.229), 1 / (255.0 * 0.224), 1 / (255.0 * 0.225)], + beta: [-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], + outInterpolation: 'lanczos' as const, +}; +const LRASPP_MOBILENET_V3_LARGE_XNNPACK_FP32: SemanticSegmentationModel = { + modelPath: `${BASE_URL}-lraspp/${VERSION_TAG}/xnnpack/lraspp_mobilenet_v3_large_xnnpack_fp32.pte`, + opts: LRASPP_MOBILENET_V3_LARGE_OPTS, +}; +const LRASPP_MOBILENET_V3_LARGE_XNNPACK_INT8: SemanticSegmentationModel = { + modelPath: `${BASE_URL}-lraspp/${VERSION_TAG}/xnnpack/lraspp_mobilenet_v3_large_xnnpack_int8.pte`, + opts: LRASPP_MOBILENET_V3_LARGE_OPTS, +}; + /** * Registry of pre-configured ExecuTorch models. * @@ -44,4 +82,15 @@ export const models = { COREML_FP16: EFFICIENTNET_V2_S_COREML_FP16, }, }, + semanticSegmentation: { + SELFIE_SEGMENTATION: { + ...SELFIE_SEGMENTATION_XNNPACK_FP32, + XNNPACK_FP32: SELFIE_SEGMENTATION_XNNPACK_FP32, + }, + LRASPP_MOBILENET_V3_LARGE: { + ...LRASPP_MOBILENET_V3_LARGE_XNNPACK_INT8, + XNNPACK_FP32: LRASPP_MOBILENET_V3_LARGE_XNNPACK_FP32, + XNNPACK_INT8: LRASPP_MOBILENET_V3_LARGE_XNNPACK_INT8, + }, + }, }; From bc774e955119462a17542ce1abaa94f783d7a728 Mon Sep 17 00:00:00 2001 From: Bartosz Hanc Date: Mon, 22 Jun 2026 16:48:36 +0200 Subject: [PATCH 2/3] refactor(cv): optimize applyColormap and document layout requirements for semantic segmentation --- .../cpp/extensions/cv/image_ops.cpp | 14 ++++++------ .../src/extensions/cv/ops/image.ts | 22 +++++++++++++------ 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp b/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp index 61ffc01704..02bf913647 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp +++ b/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp @@ -631,6 +631,9 @@ void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) { if (dst->dtype_ != rnexecutorch::core::types::DType::uint8) { throw jsi::JSError(rt, "applyColormap: dst must be uint8"); } + if (dst->numel_ != src->numel_ * 4) { + throw jsi::JSError(rt, "applyColormap: dst must have exactly 4 times the number of elements as src (RGBA channels)"); + } auto colormapArray = args[2].asObject(rt).asArray(rt); size_t numColors = colormapArray.size(rt); @@ -643,9 +646,9 @@ void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) { lut[i][3] = color.getValueAtIndex(rt, 3).asNumber(); } - std::shared_lock src_lock(src->mutex_, std::try_to_lock); - std::unique_lock dst_lock(dst->mutex_, std::try_to_lock); - if (!src_lock.owns_lock() || !dst_lock.owns_lock()) { + std::shared_lock srcLock(src->mutex_, std::try_to_lock); + std::unique_lock dstLock(dst->mutex_, std::try_to_lock); + if (!srcLock.owns_lock() || !dstLock.owns_lock()) { throw jsi::JSError(rt, "applyColormap: tensors in use"); } @@ -653,10 +656,7 @@ void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) { throw jsi::JSError(rt, "applyColormap: tensor has been disposed"); } - size_t pixels = 1; - for (int i = 0; i < src->shape_.size(); ++i) { - pixels *= src->shape_[i]; - } + size_t pixels = src->numel_; const int32_t *srcData = reinterpret_cast(src->data_.get()); uint8_t *dstData = dst->data_.get(); diff --git a/packages/react-native-executorch/src/extensions/cv/ops/image.ts b/packages/react-native-executorch/src/extensions/cv/ops/image.ts index a1fd9671d6..3fe59f7dc5 100644 --- a/packages/react-native-executorch/src/extensions/cv/ops/image.ts +++ b/packages/react-native-executorch/src/extensions/cv/ops/image.ts @@ -177,14 +177,22 @@ export function normalize(src: Tensor, dst: Tensor, opts?: NormalizeOptions): Te } /** - * Applies a colormap to a 2D single-channel image tensor, mapping class indices - * to RGBA colors. + * Applies a colormap to a single-channel image tensor, mapping class indices to + * RGBA colors. + * + * This operation iterates over each index/class ID in the source tensor, looks + * up its corresponding RGBA color in the provided colormap palette, and writes + * it to the destination tensor. * @category Typescript API - * @param src The source index/mask tensor (int32). - * @param dst The pre-allocated destination tensor (uint8, 4 channels). - * @param colormap An array of RGBA color arrays corresponding to the class - * indices. - * @returns The destination tensor with applied colormap. + * @param src The source index/mask tensor. Must be a tensor of shape `[H, W, + * 1]` (or `[H, W]`) and `int32` dtype containing class indices. + * @param dst The pre-allocated destination tensor to write the mapped RGBA + * values to. Must be a 3D image tensor in HWC layout of shape `[H, W, 4]` and + * `uint8` dtype. + * @param colormap An array of RGBA color arrays `[R, G, B, A]` corresponding to + * each class index. The size of this list must cover all class indices present + * in `src`. + * @returns The destination tensor with the applied colormap. */ export function applyColormap( src: Tensor, From ff18eb802e070f033d0a4204dd897be6ccb0c305 Mon Sep 17 00:00:00 2001 From: Bartosz Hanc Date: Mon, 22 Jun 2026 17:12:32 +0200 Subject: [PATCH 3/3] refactor(cv): validate colormap input parameters and document layout formats --- .../cpp/extensions/cv/image_ops.cpp | 31 ++++++++++++++++--- .../src/extensions/cv/ops/image.ts | 13 ++++---- 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp b/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp index 02bf913647..c8c3861b04 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp +++ b/packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp @@ -622,6 +622,13 @@ void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) { throw jsi::JSError(rt, "Usage: applyColormap(src, dst, colormap)"); } + if (!args[0].isObject() || !args[0].asObject(rt).isHostObject(rt)) { + throw jsi::JSError(rt, "applyColormap: src must be a Tensor"); + } + if (!args[1].isObject() || !args[1].asObject(rt).isHostObject(rt)) { + throw jsi::JSError(rt, "applyColormap: dst must be a Tensor"); + } + auto src = args[0].asObject(rt).getHostObject(rt); auto dst = args[1].asObject(rt).getHostObject(rt); @@ -635,15 +642,29 @@ void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) { 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> lut(numColors); for (size_t i = 0; i < numColors; ++i) { - auto color = colormapArray.getValueAtIndex(rt, i).asObject(rt).asArray(rt); - lut[i][0] = color.getValueAtIndex(rt, 0).asNumber(); - lut[i][1] = color.getValueAtIndex(rt, 1).asNumber(); - lut[i][2] = color.getValueAtIndex(rt, 2).asNumber(); - lut[i][3] = color.getValueAtIndex(rt, 3).asNumber(); + 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) != 4) { + throw jsi::JSError(rt, "applyColormap: colormap entry must be an RGBA color array of size 4"); + } + for (size_t c = 0; c < 4; ++c) { + auto channelVal = color.getValueAtIndex(rt, c); + if (!channelVal.isNumber()) { + throw jsi::JSError(rt, "applyColormap: colormap channel value must be a number"); + } + lut[i][c] = static_cast(channelVal.asNumber()); + } } std::shared_lock srcLock(src->mutex_, std::try_to_lock); diff --git a/packages/react-native-executorch/src/extensions/cv/ops/image.ts b/packages/react-native-executorch/src/extensions/cv/ops/image.ts index 3fe59f7dc5..faffb733d1 100644 --- a/packages/react-native-executorch/src/extensions/cv/ops/image.ts +++ b/packages/react-native-executorch/src/extensions/cv/ops/image.ts @@ -184,14 +184,13 @@ export function normalize(src: Tensor, dst: Tensor, opts?: NormalizeOptions): Te * up its corresponding RGBA color in the provided colormap palette, and writes * it to the destination tensor. * @category Typescript API - * @param src The source index/mask tensor. Must be a tensor of shape `[H, W, - * 1]` (or `[H, W]`) and `int32` dtype containing class indices. + * @param src The source index/mask tensor. Must be an integer tensor of `int32` + * dtype containing class indices. Shape `[H, W, 1]` (or `[H, W]`). * @param dst The pre-allocated destination tensor to write the mapped RGBA - * values to. Must be a 3D image tensor in HWC layout of shape `[H, W, 4]` and - * `uint8` dtype. - * @param colormap An array of RGBA color arrays `[R, G, B, A]` corresponding to - * each class index. The size of this list must cover all class indices present - * in `src`. + * values to. Must be a 3D image tensor in HWC layout and `uint8` dtype. Shape + * `[H, W, 4]`. + * @param colormap An array of RGBA color arrays `[R, G, B, A]` corresponding to each + * class index. The size of this list must cover all class indices present in `src`. * @returns The destination tensor with the applied colormap. */ export function applyColormap(