From e968b5cd0393d32343c69fc77c294b07f658b03e Mon Sep 17 00:00:00 2001 From: Bartosz Hanc Date: Fri, 26 Jun 2026 21:52:59 +0200 Subject: [PATCH 1/2] feat: add instance segmentation pipeline --- .../cpp/extensions/cv/box_ops.cpp | 134 ++++++++++++ .../cpp/extensions/cv/box_ops.h | 1 + .../cpp/extensions/cv/install.cpp | 1 + .../cpp/extensions/math/install.cpp | 1 + .../cpp/extensions/math/operations.cpp | 67 ++++++ .../cpp/extensions/math/operations.h | 1 + .../src/extensions/cv/image.ts | 2 +- .../src/extensions/cv/ops/boxes.ts | 31 +++ .../src/extensions/cv/ops/image.ts | 10 +- .../cv/tasks/instanceSegmentation.ts | 197 ++++++++++++++++++ .../src/extensions/math.ts | 15 ++ .../src/hooks/useInstanceSegmenter.ts | 49 +++++ packages/react-native-executorch/src/index.ts | 2 + 13 files changed, 506 insertions(+), 5 deletions(-) create mode 100644 packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts create mode 100644 packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts diff --git a/packages/react-native-executorch/cpp/extensions/cv/box_ops.cpp b/packages/react-native-executorch/cpp/extensions/cv/box_ops.cpp index 5d42fc5097..28cbf7eb98 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/box_ops.cpp +++ b/packages/react-native-executorch/cpp/extensions/cv/box_ops.cpp @@ -9,6 +9,8 @@ #include #include +#include + #include "core/dtype.h" #include "core/tensor.h" @@ -243,4 +245,136 @@ void install_nms(jsi::Runtime &rt, jsi::Object &module) { module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 3, fnBody)); } + +namespace { +int dtypeToCvDepth(rnexecutorch::core::types::DType dtype) { + switch (dtype) { + case rnexecutorch::core::types::DType::uint8: + return CV_8U; + case rnexecutorch::core::types::DType::int32: + return CV_32S; + case rnexecutorch::core::types::DType::float32: + return CV_32F; + } + throw std::invalid_argument("unsupported dtype"); +} +} // namespace + +void install_restrictToBox(jsi::Runtime &rt, jsi::Object &module) { + auto name = "restrictToBox"; + auto fnBody = [](jsi::Runtime &rt, const jsi::Value & /*thisVal*/, const jsi::Value *args, size_t count) -> jsi::Value { + if (count != 4) { + throw jsi::JSError(rt, "Usage: restrictToBox(src, dst, boxTuple, format)"); + } + + if (!args[0].isObject() || !args[0].asObject(rt).isHostObject(rt)) { + throw jsi::JSError(rt, "restrictToBox: src must be a Tensor"); + } + + if (!args[1].isObject() || !args[1].asObject(rt).isHostObject(rt)) { + throw jsi::JSError(rt, "restrictToBox: dst must be a Tensor"); + } + + if (!args[2].isObject() || !args[2].asObject(rt).isArray(rt)) { + throw jsi::JSError(rt, "restrictToBox: boxTuple must be an Array"); + } + + if (!args[3].isString()) { + throw jsi::JSError(rt, "restrictToBox: format must be a String"); + } + + auto src = args[0].asObject(rt).getHostObject(rt); + auto dst = args[1].asObject(rt).getHostObject(rt); + auto boxTuple = args[2].asObject(rt).asArray(rt); + std::string boxFormatStr = args[3].asString(rt).utf8(rt); + + if (boxTuple.size(rt) != 4) { + throw jsi::JSError(rt, "restrictToBox: boxTuple must contain exactly 4 coordinates"); + } + + BoxFormat boxFormat; + try { + boxFormat = parseBoxFormat(boxFormatStr); + } catch (const std::invalid_argument &e) { + throw jsi::JSError(rt, "restrictToBox: " + std::string(e.what())); + } + + float a = static_cast(boxTuple.getValueAtIndex(rt, 0).asNumber()); + float b = static_cast(boxTuple.getValueAtIndex(rt, 1).asNumber()); + float c = static_cast(boxTuple.getValueAtIndex(rt, 2).asNumber()); + float d = static_cast(boxTuple.getValueAtIndex(rt, 3).asNumber()); + + auto [xmin, ymin, xmax, ymax] = decodeToXyxy(a, b, c, d, boxFormat); + + if (src.get() == dst.get()) { + throw jsi::JSError(rt, "restrictToBox: In-place operations (src == dst) are not supported."); + } + + if (src->shape_ != dst->shape_) { + throw jsi::JSError(rt, "restrictToBox: src and dst must have the same shape"); + } + + if (src->shape_.size() < 2) { + throw jsi::JSError(rt, "restrictToBox: src must have at least 2 dimensions [H, W, ...]"); + } + + int32_t H = src->shape_[0]; + int32_t W = src->shape_[1]; + int32_t C = 1; + for (size_t i = 2; i < src->shape_.size(); ++i) { + C *= src->shape_[i]; + } + + int32_t x1 = static_cast(std::ceil(xmin)); + int32_t y1 = static_cast(std::ceil(ymin)); + int32_t x2 = static_cast(std::floor(xmax)); + int32_t y2 = static_cast(std::floor(ymax)); + + x1 = std::max(0, x1); + y1 = std::max(0, y1); + x2 = std::min(W - 1, x2); + y2 = std::min(H - 1, y2); + + bool isEmpty = (x2 < x1) || (y2 < y1); + + if (src->dtype_ != dst->dtype_) { + throw jsi::JSError(rt, "restrictToBox: src and dst must have the same dtype"); + } + + 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, "restrictToBox: tensors in use"); + } + + if (!src->data_ || !dst->data_) { + throw jsi::JSError(rt, "restrictToBox: tensors must not be disposed"); + } + + int32_t cvType; + try { + cvType = CV_MAKETYPE(dtypeToCvDepth(src->dtype_), C); + } catch (const std::invalid_argument &e) { + throw jsi::JSError(rt, "restrictToBox: " + std::string(e.what())); + } + + ::cv::Mat srcMat(H, W, cvType, src->data_.get()); + ::cv::Mat dstMat(H, W, cvType, dst->data_.get()); + + if (isEmpty) { + dstMat.setTo(::cv::Scalar::all(0)); + } else { + dstMat.setTo(::cv::Scalar::all(0)); + int32_t boxW = x2 - x1 + 1; + int32_t boxH = y2 - y1 + 1; + ::cv::Rect roi(x1, y1, boxW, boxH); + srcMat(roi).copyTo(dstMat(roi)); + } + + return jsi::Value(rt, args[1]); + }; + + module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 6, fnBody)); +} + } // namespace rnexecutorch::extensions::cv::box_ops diff --git a/packages/react-native-executorch/cpp/extensions/cv/box_ops.h b/packages/react-native-executorch/cpp/extensions/cv/box_ops.h index 7e001b67b7..6bc14c772f 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/box_ops.h +++ b/packages/react-native-executorch/cpp/extensions/cv/box_ops.h @@ -4,4 +4,5 @@ namespace rnexecutorch::extensions::cv::box_ops { void install_nms(facebook::jsi::Runtime &rt, facebook::jsi::Object &module); +void install_restrictToBox(facebook::jsi::Runtime &rt, facebook::jsi::Object &module); } // namespace rnexecutorch::extensions::cv::box_ops diff --git a/packages/react-native-executorch/cpp/extensions/cv/install.cpp b/packages/react-native-executorch/cpp/extensions/cv/install.cpp index ae370a9046..0540f45c9b 100644 --- a/packages/react-native-executorch/cpp/extensions/cv/install.cpp +++ b/packages/react-native-executorch/cpp/extensions/cv/install.cpp @@ -16,6 +16,7 @@ void install(facebook::jsi::Runtime &rt, facebook::jsi::Object &module) { image_ops::install_applyColormap(rt, cvModule); box_ops::install_nms(rt, cvModule); + box_ops::install_restrictToBox(rt, cvModule); module.setProperty(rt, "cv", cvModule); } diff --git a/packages/react-native-executorch/cpp/extensions/math/install.cpp b/packages/react-native-executorch/cpp/extensions/math/install.cpp index 39bce39544..7a7e3073e6 100644 --- a/packages/react-native-executorch/cpp/extensions/math/install.cpp +++ b/packages/react-native-executorch/cpp/extensions/math/install.cpp @@ -10,6 +10,7 @@ void install(jsi::Runtime &rt, jsi::Object &module) { install_sigmoid(rt, mathModule); install_softmax(rt, mathModule); install_argmax(rt, mathModule); + install_threshold(rt, mathModule); module.setProperty(rt, "math", mathModule); } diff --git a/packages/react-native-executorch/cpp/extensions/math/operations.cpp b/packages/react-native-executorch/cpp/extensions/math/operations.cpp index 4d6fdde5ff..5483182da8 100644 --- a/packages/react-native-executorch/cpp/extensions/math/operations.cpp +++ b/packages/react-native-executorch/cpp/extensions/math/operations.cpp @@ -301,4 +301,71 @@ void install_argmax(jsi::Runtime &rt, jsi::Object &module) { }; module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 3, fnBody)); } + +void install_threshold(jsi::Runtime &rt, jsi::Object &module) { + auto name = "threshold"; + 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: threshold(src, dst, threshold)"); + } + + if (!args[0].isObject() || !args[0].asObject(rt).isHostObject(rt)) { + throw jsi::JSError(rt, "threshold: src must be a Tensor"); + } + + if (!args[1].isObject() || !args[1].asObject(rt).isHostObject(rt)) { + throw jsi::JSError(rt, "threshold: dst must be a Tensor"); + } + + if (!args[2].isNumber()) { + throw jsi::JSError(rt, "threshold: threshold must be a number"); + } + + auto src = args[0].asObject(rt).getHostObject(rt); + auto dst = args[1].asObject(rt).getHostObject(rt); + float thresholdVal = static_cast(args[2].asNumber()); + + if (src.get() == dst.get()) { + throw jsi::JSError(rt, "threshold: In-place operations (src == dst) are not supported."); + } + + if (src->shape_ != dst->shape_) { + throw jsi::JSError(rt, "threshold: src and dst must have the same shape"); + } + + if (src->dtype_ != rnexecutorch::core::types::DType::float32) { + throw jsi::JSError(rt, "threshold: src must be a float32 tensor"); + } + + if (dst->dtype_ != rnexecutorch::core::types::DType::float32) { + throw jsi::JSError(rt, "threshold: dst must be a float32 tensor"); + } + + 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, "threshold: tensors in use"); + } + + if (!src->data_) { + throw jsi::JSError(rt, "threshold: src tensor has been disposed"); + } + + if (!dst->data_) { + throw jsi::JSError(rt, "threshold: dst tensor has been disposed"); + } + + const auto *srcData = reinterpret_cast(src->data_.get()); + auto *dstData = reinterpret_cast(dst->data_.get()); + + for (size_t i = 0; i < src->numel_; ++i) { + dstData[i] = (srcData[i] >= thresholdVal) ? 1.0f : 0.0f; + } + + return jsi::Value(rt, args[1]); + }; + + module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 3, fnBody)); +} + } // namespace rnexecutorch::extensions::math diff --git a/packages/react-native-executorch/cpp/extensions/math/operations.h b/packages/react-native-executorch/cpp/extensions/math/operations.h index 03455c30f8..9196a036ec 100644 --- a/packages/react-native-executorch/cpp/extensions/math/operations.h +++ b/packages/react-native-executorch/cpp/extensions/math/operations.h @@ -6,4 +6,5 @@ namespace rnexecutorch::extensions::math { void install_sigmoid(facebook::jsi::Runtime &rt, facebook::jsi::Object &module); void install_softmax(facebook::jsi::Runtime &rt, facebook::jsi::Object &module); void install_argmax(facebook::jsi::Runtime &rt, facebook::jsi::Object &module); +void install_threshold(facebook::jsi::Runtime &rt, facebook::jsi::Object &module); } // namespace rnexecutorch::extensions::math diff --git a/packages/react-native-executorch/src/extensions/cv/image.ts b/packages/react-native-executorch/src/extensions/cv/image.ts index c217f0ae29..54856b02e4 100644 --- a/packages/react-native-executorch/src/extensions/cv/image.ts +++ b/packages/react-native-executorch/src/extensions/cv/image.ts @@ -2,7 +2,7 @@ * Supported pixel format layouts for image buffers. * @category Types */ -export type ImageFormat = 'rgb' | 'rgba' | 'bgr' | 'bgra'; +export type ImageFormat = 'rgb' | 'rgba' | 'bgr' | 'bgra' | 'gray'; /** * Represents a raw CPU image buffer in HWC (Height, Width, Channel) layout. diff --git a/packages/react-native-executorch/src/extensions/cv/ops/boxes.ts b/packages/react-native-executorch/src/extensions/cv/ops/boxes.ts index 90aaeefeb1..0d1b8e2bf9 100644 --- a/packages/react-native-executorch/src/extensions/cv/ops/boxes.ts +++ b/packages/react-native-executorch/src/extensions/cv/ops/boxes.ts @@ -165,3 +165,34 @@ export function nms(boxes: Tensor, scores: Tensor, opts: NmsOptions): number[] | 'worklet'; return rnexecutorchJsi.cv.nms(boxes, scores, opts); } + +/** + * Masks the source tensor by keeping only the elements inside the specified + * bounding box, writing the result to a pre-allocated destination tensor. + * + * Note: This operation does not change the tensor dimensions (it does not crop + * the shape). Instead, it copies the elements within the box coordinates from + * `src` to `dst`, and sets all elements outside the box to `0`. + * @category Typescript API + * @param src The source tensor of shape [H, W, ...] with at least 2 dimensions. + * @param dst The pre-allocated destination tensor of the same shape and data + * type as `src`. + * @param box The bounding box defining the region of interest to copy. + * @returns The destination tensor containing the masked output. + */ +export function restrictToBox(src: Tensor, dst: Tensor, box: BoundingBox): Tensor { + 'worklet'; + let [a, b, c, d] = [0, 0, 0, 0]; + switch (box.format) { + case 'xyxy': + [a, b, c, d] = [box.xmin, box.ymin, box.xmax, box.ymax]; + break; + case 'xywh': + [a, b, c, d] = [box.xmin, box.ymin, box.w, box.h]; + break; + case 'cxcywh': + [a, b, c, d] = [box.cx, box.cy, box.w, box.h]; + break; + } + return rnexecutorchJsi.cv.restrictToBox(src, dst, [a, b, c, d], box.format); +} 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 b5bc22b746..dc930eceec 100644 --- a/packages/react-native-executorch/src/extensions/cv/ops/image.ts +++ b/packages/react-native-executorch/src/extensions/cv/ops/image.ts @@ -31,10 +31,11 @@ export const FORMAT_CONVERSION: Record< ImageFormat, Record > = { - rgb: { rgb: null, rgba: 'RGB2RGBA', bgr: 'RGB2BGR', bgra: null }, - bgr: { rgb: 'BGR2RGB', rgba: 'BGR2RGBA', bgr: null, bgra: null }, - rgba: { rgb: 'RGBA2RGB', rgba: null, bgr: 'RGBA2BGR', bgra: null }, - bgra: { rgb: 'BGRA2RGB', rgba: 'BGRA2RGBA', bgr: 'BGRA2BGR', bgra: null }, + rgb: { rgb: null, rgba: 'RGB2RGBA', bgr: 'RGB2BGR', bgra: null, gray: 'RGB2GRAY' }, + bgr: { rgb: 'BGR2RGB', rgba: 'BGR2RGBA', bgr: null, bgra: null, gray: 'BGR2GRAY' }, + rgba: { rgb: 'RGBA2RGB', rgba: null, bgr: 'RGBA2BGR', bgra: null, gray: 'RGBA2GRAY' }, + bgra: { rgb: 'BGRA2RGB', rgba: 'BGRA2RGBA', bgr: 'BGRA2BGR', bgra: null, gray: 'BGRA2GRAY' }, + gray: { rgb: 'GRAY2RGBA', rgba: 'GRAY2RGBA', bgr: null, bgra: null, gray: null }, }; /** @@ -47,6 +48,7 @@ export const FORMAT_CHANNELS: Record = { bgr: 3, rgba: 4, bgra: 4, + gray: 1, }; /** diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts new file mode 100644 index 0000000000..c9c40653eb --- /dev/null +++ b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts @@ -0,0 +1,197 @@ +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 { threshold } from '../../math'; +import { resize, normalize } from '../ops/image'; +import { + decodeBox, + scaleBox, + nms, + restrictToBox, + type BoundingBox, + type BoxFormat, +} from '../ops/boxes'; + +export type { BoxFormat }; + +export type InstanceSegmenterOptions = Omit< + ImagePreprocessorOptions, + 'resizeMode' +> & { + readonly resizeMode: 'stretch'; + readonly labels: readonly L[]; + readonly boxFormat: F; + readonly defaultIouThreshold: number; + readonly defaultMaskThreshold: number; + readonly defaultConfidenceThreshold: number; +}; + +export type InstanceSegmenterModel = { + readonly modelPath: string; + readonly opts: InstanceSegmenterOptions; +}; + +export type InstanceSegmentation = { + readonly box: BoundingBox; + readonly mask: ImageBuffer; + readonly label: L; + readonly confidence: number; +}; + +export async function createInstanceSegmenter( + config: InstanceSegmenterModel, + runtime?: WorkletRuntime +): Promise<{ + dispose: () => void; + + segmentInstances: ( + input: ImageBuffer, + options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } + ) => Promise[]>; + + segmentInstancesWorklet: ( + input: ImageBuffer, + options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } + ) => InstanceSegmentation[]; +}> { + 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', ['N', 4]), + SymbolicTensor('float32', ['N']), + SymbolicTensor('float32', ['N']), + SymbolicTensor('float32', ['N', 'MH', 'MW']), + ] + ); + + const inpShape = meta.inputTensorMeta[0]!.shape; + + const outBoxesShape = meta.outputTensorMeta[0]!.shape; + const outScoresShape = meta.outputTensorMeta[1]!.shape; + const outClassesShape = meta.outputTensorMeta[2]!.shape; + const outMasksShape = meta.outputTensorMeta[3]!.shape; + + const maskH = outMasksShape[1]!; + const maskW = outMasksShape[2]!; + const targetH = inpShape.at(-2)!; + const targetW = inpShape.at(-1)!; + + const tensors = [ + tensor('float32', outBoxesShape), + tensor('float32', outScoresShape), + tensor('float32', outClassesShape), + tensor('float32', outMasksShape), + tensor('float32', [maskH, maskW, 1]), + ] as const; + + const [tBoxes, tScores, tClasses, tAllMasks, tMask] = tensors; + + const preprocessor = createImagePreprocessor(opts, inpShape); + + const dispose = () => { + preprocessor.dispose(); + tensors.forEach((t) => t.dispose()); + model.dispose(); + }; + + const segmentInstancesWorklet = ( + input: ImageBuffer, + options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } + ): InstanceSegmentation[] => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tBoxes, tScores, tClasses, tAllMasks]); + + const iouThreshold = options?.iouThreshold ?? opts.defaultIouThreshold; + const maskThreshold = options?.maskThreshold ?? opts.defaultMaskThreshold; + const confidenceThreshold = options?.confidenceThreshold ?? opts.defaultConfidenceThreshold; + + const eps = 1e-7; + const clampedMaskThreshold = Math.max(eps, Math.min(1 - eps, maskThreshold)); + const logitMaskThreshold = Math.log(clampedMaskThreshold / (1 - clampedMaskThreshold)); + + const indices = nms(tBoxes, tScores, { + boxFormat: opts.boxFormat, + iouThreshold, + confidenceThreshold, + nmsType: 'standard', + }); + + const boxes = tBoxes.getData(new Float32Array(tBoxes.numel)); + const scores = tScores.getData(new Float32Array(tScores.numel)); + const classes = tClasses.getData(new Float32Array(tClasses.numel)); + + const auxTensors = [ + tensor('float32', [input.height, input.width, 1]), + tensor('float32', [input.height, input.width, 1]), + tensor('float32', [input.height, input.width, 1]), + tensor('uint8', [input.height, input.width, 1]), + ] as const; + + const [tResize, tThreshold, tCrop, tUint8] = auxTensors; + + const results: InstanceSegmentation[] = []; + + try { + for (const idx of indices) { + const confidence = scores[idx]!; + const classIdx = Math.round(classes[idx]!); + const label = opts.labels[classIdx]; + + if (label === undefined) { + throw new Error( + `InstanceSegmenter: Predicted class index ${classIdx} is` + + `out of bounds for labels array of size ${opts.labels.length}.` + ); + } + + const a = boxes[idx * 4]!; + const b = boxes[idx * 4 + 1]!; + const c = boxes[idx * 4 + 2]!; + const d = boxes[idx * 4 + 3]!; + + const box = scaleBox(decodeBox([a, b, c, d], opts.boxFormat), { + from: { width: targetW, height: targetH }, + to: { width: input.width, height: input.height }, + resizeMode: 'stretch', + }); + + const maskData = tAllMasks + .copyTo(tMask, { offset: idx * maskH * maskW, length: maskH * maskW }) + .through(resize, tResize, { mode: 'stretch', interpolation: 'linear' }) + .through(threshold, tThreshold, logitMaskThreshold) + .through(restrictToBox, tCrop, box) + .through(normalize, tUint8, { alpha: 255.0 }) + .getData(new Uint8Array(tUint8.numel)); + + const mask = { + data: maskData, + width: input.width, + height: input.height, + format: 'gray' as const, + layout: 'hwc' as const, + }; + + results.push({ box, mask, confidence, label }); + } + } finally { + auxTensors.forEach((t) => t.dispose()); + } + + return results; + }; + + const segmentInstances = wrapAsync(segmentInstancesWorklet, runtime); + + return { segmentInstances, segmentInstancesWorklet, dispose }; +} diff --git a/packages/react-native-executorch/src/extensions/math.ts b/packages/react-native-executorch/src/extensions/math.ts index 679a8567db..11f78ea080 100644 --- a/packages/react-native-executorch/src/extensions/math.ts +++ b/packages/react-native-executorch/src/extensions/math.ts @@ -46,3 +46,18 @@ export function argmax(src: Tensor, dst: Tensor, axis: number = -1): Tensor { 'worklet'; return rnexecutorchJsi.math.argmax(src, dst, axis); } + +/** + * Applies the element-wise threshold step function on a float32 source tensor and + * writes the result to a destination tensor. + * @category Typescript API + * @param src The input float32 source tensor. Shape [d1,...,dn]. + * @param dst The pre-allocated destination tensor to write the result to. + * `dst` tensor must have the same shape as `src` and have dtype float32. + * @param thresholdVal The threshold value above or equal to which elements are mapped to 1.0. + * @returns The destination tensor containing the threshold step output. + */ +export function threshold(src: Tensor, dst: Tensor, thresholdVal: number): Tensor { + 'worklet'; + return rnexecutorchJsi.math.threshold(src, dst, thresholdVal); +} diff --git a/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts b/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts new file mode 100644 index 0000000000..dd82bc95eb --- /dev/null +++ b/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts @@ -0,0 +1,49 @@ +import { useModel } from './useModel'; +import { useResourceDownload } from './useResourceDownload'; +import { + createInstanceSegmenter, + type InstanceSegmenterModel, +} from '../extensions/cv/tasks/instanceSegmentation'; +import type { BoxFormat } from '../extensions/cv/ops/boxes'; + +/** + * React hook to load and run an instance 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 F The bounding box format. + * @typeParam L The class labels type. + * @param config The instance 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 useInstanceSegmenter( + config: InstanceSegmenterModel, + options?: { preventLoad?: boolean } +) { + const { localPath, downloadProgress, downloadError } = useResourceDownload( + config.modelPath, + options?.preventLoad + ); + const { model, error } = useModel( + createInstanceSegmenter, + localPath ? { ...config, modelPath: localPath } : null, + [localPath] + ); + + return { + isReady: !!model, + error: downloadError || error, + downloadProgress, + localPath, + segmentInstances: model?.segmentInstances, + segmentInstancesWorklet: model?.segmentInstancesWorklet, + labels: config.opts.labels, + }; +} diff --git a/packages/react-native-executorch/src/index.ts b/packages/react-native-executorch/src/index.ts index 47fbd97b87..00557be78b 100644 --- a/packages/react-native-executorch/src/index.ts +++ b/packages/react-native-executorch/src/index.ts @@ -2,6 +2,7 @@ export * from './hooks/useClassifier'; export * from './hooks/useStyleTransfer'; export * from './hooks/useSemanticSegmenter'; +export * from './hooks/useInstanceSegmenter'; export * from './hooks/useKeypointDetector'; export * from './hooks/useObjectDetector'; export * from './hooks/useTokenizer'; @@ -16,6 +17,7 @@ export * as constants from './constants'; export * from './extensions/cv/tasks/classification'; export * from './extensions/cv/tasks/styleTransfer'; export * from './extensions/cv/tasks/semanticSegmentation'; +export * from './extensions/cv/tasks/instanceSegmentation'; export * from './extensions/cv/tasks/keypointDetection'; export * from './extensions/cv/tasks/objectDetection'; export * from './extensions/nlp/tasks/tokenization'; From b0d203d52fe6008e1f8e07d33810175ba05560c1 Mon Sep 17 00:00:00 2001 From: Bartosz Hanc Date: Fri, 26 Jun 2026 22:42:55 +0200 Subject: [PATCH 2/2] style: small fixes --- .../extensions/cv/tasks/instanceSegmentation.ts | 14 ++++++++------ .../src/hooks/useInstanceSegmenter.ts | 2 +- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts index c9c40653eb..84eea82a7e 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts @@ -37,13 +37,15 @@ export type InstanceSegmenterModel = { readonly opts: InstanceSegmenterOptions; }; -export type InstanceSegmentation = { +export type InstanceSegmentationResult = { readonly box: BoundingBox; readonly mask: ImageBuffer; readonly label: L; readonly confidence: number; }; +export type InstanceSegmentation = InstanceSegmentationResult; + export async function createInstanceSegmenter( config: InstanceSegmenterModel, runtime?: WorkletRuntime @@ -53,12 +55,12 @@ export async function createInstanceSegmenter( segmentInstances: ( input: ImageBuffer, options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } - ) => Promise[]>; + ) => Promise[]>; segmentInstancesWorklet: ( input: ImageBuffer, options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } - ) => InstanceSegmentation[]; + ) => InstanceSegmentationResult[]; }> { const { modelPath, opts } = config; const model = await wrapAsync(loadModel, runtime)(modelPath); @@ -107,7 +109,7 @@ export async function createInstanceSegmenter( const segmentInstancesWorklet = ( input: ImageBuffer, options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } - ): InstanceSegmentation[] => { + ): InstanceSegmentationResult[] => { 'worklet'; const tInput = preprocessor.process(input); model.execute('forward', [tInput], [tBoxes, tScores, tClasses, tAllMasks]); @@ -140,7 +142,7 @@ export async function createInstanceSegmenter( const [tResize, tThreshold, tCrop, tUint8] = auxTensors; - const results: InstanceSegmentation[] = []; + const results: InstanceSegmentationResult[] = []; try { for (const idx of indices) { @@ -150,7 +152,7 @@ export async function createInstanceSegmenter( if (label === undefined) { throw new Error( - `InstanceSegmenter: Predicted class index ${classIdx} is` + + `InstanceSegmenter: Predicted class index ${classIdx} is ` + `out of bounds for labels array of size ${opts.labels.length}.` ); } diff --git a/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts b/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts index dd82bc95eb..364f2b32a7 100644 --- a/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts +++ b/packages/react-native-executorch/src/hooks/useInstanceSegmenter.ts @@ -3,8 +3,8 @@ import { useResourceDownload } from './useResourceDownload'; import { createInstanceSegmenter, type InstanceSegmenterModel, + type BoxFormat, } from '../extensions/cv/tasks/instanceSegmentation'; -import type { BoxFormat } from '../extensions/cv/ops/boxes'; /** * React hook to load and run an instance segmentation model.