feat: make classification module a vision labeled module

Author: NorbertKlockiewicz

.cspell-wordlist.txt

@@ -129,6 +129,45 @@ metaprogramming
 ktlint
 lefthook
 espeak
+KOMODO
+DUNGENESS
+SHIH
+RIDGEBACK
+BLUETICK
+REDBONE
+IBIZAN
+OTTERHOUND
+BULLTERRIER
+BEDLINGTON
+SEALYHAM
+DANDIE
+DINMONT
+VIZSLA
+CLUMBER
+MALINOIS
+KOMONDOR
+BOUVIER
+FLANDRES
+APPENZELLER
+ENTLEBUCHER
+LEONBERG
+BRABANCON
+LYCAENID
+PATAS
+INDRI
+BARRACOUTA
+ABAYA
+BOTTLECAP
+CHAINLINK
+GASMASK
+GOLFCART
+HOOPSKIRT
+LUMBERMILL
+PADDLEWHEEL
+PICKELHAUBE
+CARBONARA
+GYROMITRA
+BOLETE
 NCHW
 həlˈO
 wˈɜɹld

docs/docs/03-hooks/02-computer-vision/useClassification.md

@@ -37,9 +37,13 @@ try {
 
 `useClassification` takes [`ClassificationProps`](../../06-api-reference/interfaces/ClassificationProps.md) that consists of:
 
-- `model` containing [`modelSource`](../../06-api-reference/interfaces/ClassificationProps.md#modelsource).
+- `model` - An object containing:
+  - `modelName` - The name of a built-in model. See [`ClassificationModelSources`](../../06-api-reference/interfaces/ClassificationProps.md) for the list of supported models.
+  - `modelSource` - The location of the model binary (a URL or a bundled resource).
 - An optional flag [`preventLoad`](../../06-api-reference/interfaces/ClassificationProps.md#preventload) which prevents auto-loading of the model.
 
+The hook is generic over the model config — TypeScript automatically infers the correct label type based on the `modelName` you provide. No explicit generic parameter is needed.
+
 You need more details? Check the following resources:
 
 - For detailed information about `useClassification` arguments check this section: [`useClassification` arguments](../../06-api-reference/functions/useClassification.md#parameters).
@@ -48,11 +52,17 @@ You need more details? Check the following resources:
 
 ### Returns
 
-`useClassification` returns an object called `ClassificationType` containing bunch of functions to interact with Classification models. To get more details please read: [`ClassificationType` API Reference](../../06-api-reference/interfaces/ClassificationType.md).
+`useClassification` returns a [`ClassificationType`](../../06-api-reference/interfaces/ClassificationType.md) object containing:
+
+- `isReady` - Whether the model is loaded and ready to process images.
+- `isGenerating` - Whether the model is currently processing an image.
+- `error` - An error object if the model failed to load or encountered a runtime error.
+- `downloadProgress` - A value between 0 and 1 representing the download progress of the model binary.
+- `forward` - A function to run inference on an image.
 
 ## Running the model
 
-To run the model, use the [`forward`](../../06-api-reference/interfaces/ClassificationType.md#forward) method. It accepts one argument — the image to classify. The image can be a remote URL, a local file URI, a base64-encoded image (whole URI or only raw base64), or a [`PixelData`](../../06-api-reference/interfaces/PixelData.md) object (raw RGB pixel buffer). The function returns a promise resolving to an object containing categories with their probabilities.
+To run the model, use the [`forward`](../../06-api-reference/interfaces/ClassificationType.md#forward) method. It accepts one argument — the image to classify. The image can be a remote URL, a local file URI, a base64-encoded image (whole URI or only raw base64), or a [`PixelData`](../../06-api-reference/interfaces/PixelData.md) object (raw RGB pixel buffer). The function returns a promise resolving to an object mapping label keys to their probabilities.
 
 :::info
 Images from external sources are stored in your application's temporary directory.

docs/docs/04-typescript-api/02-computer-vision/ClassificationModule.md

@@ -33,24 +33,45 @@ All methods of `ClassificationModule` are explained in details here: [`Classific
 
 ## Loading the model
 
-To create a ready-to-use instance, call the static [`fromModelName`](../../06-api-reference/classes/ClassificationModule.md#frommodelname) factory with the following parameters:
-
-- `namedSources` - Object containing:
-  - `modelName` - Model name identifier.
-  - `modelSource` - Location of the model binary.
-
-- `onDownloadProgress` - Optional callback to track download progress (value between 0 and 1).
-
-The factory returns a promise that resolves to a loaded `ClassificationModule` instance.
+Use the static [`fromModelName`](../../06-api-reference/classes/ClassificationModule.md#frommodelname) factory method. It accepts a model config object (e.g. `EFFICIENTNET_V2_S`) and an optional `onDownloadProgress` callback. It returns a promise resolving to a `ClassificationModule` instance.
 
 For more information on loading resources, take a look at [loading models](../../01-fundamentals/02-loading-models.md) page.
 
 ## Running the model
 
-To run the model, use the [`forward`](../../06-api-reference/classes/ClassificationModule.md#forward) method. It accepts one argument — the image to classify. The image can be a remote URL, a local file URI, a base64-encoded image (whole URI or only raw base64), or a [`PixelData`](../../06-api-reference/interfaces/PixelData.md) object (raw RGB pixel buffer). The method returns a promise resolving to an object containing categories with their probabilities.
+To run the model, use the [`forward`](../../06-api-reference/classes/ClassificationModule.md#forward) method. It accepts one argument — the image to classify. The image can be a remote URL, a local file URI, a base64-encoded image (whole URI or only raw base64), or a [`PixelData`](../../06-api-reference/interfaces/PixelData.md) object (raw RGB pixel buffer). The method returns a promise resolving to an object mapping label keys to their probabilities.
 
 For real-time frame processing, use [`runOnFrame`](../../03-hooks/02-computer-vision/visioncamera-integration.md) instead.
 
+## Using a custom model
+
+Use [`fromCustomModel`](../../06-api-reference/classes/ClassificationModule.md#fromcustommodel) to load your own exported model binary instead of a built-in preset.
+
+```typescript
+import { ClassificationModule } from 'react-native-executorch';
+
+const MyLabels = { CAT: 0, DOG: 1, BIRD: 2 } as const;
+
+const classifier = await ClassificationModule.fromCustomModel(
+  'https://example.com/custom_classifier.pte',
+  { labelMap: MyLabels },
+  (progress) => console.log(progress)
+);
+
+const result = await classifier.forward(imageUri);
+// result is typed as Record<'CAT' | 'DOG' | 'BIRD', number>
+```
+
+### Required model contract
+
+The `.pte` binary must expose a single `forward` method with the following interface:
+
+**Input:** one `float32` tensor of shape `[1, 3, H, W]` — a single RGB image, values in `[0, 1]` after optional per-channel normalization `(pixel − mean) / std`. H and W are read from the model's declared input shape at load time.
+
+**Output:** one `float32` tensor of shape `[1, C]` containing raw logits — one value per class, in the same order as the entries in your `labelMap`. Softmax is applied by the native runtime.
+
+Preprocessing (resize → normalize) is handled by the native runtime — your model only needs to produce the raw logits.
+
 ## Managing memory
 
-The module is a regular JavaScript object, and as such its lifespan will be managed by the garbage collector. In most cases this should be enough, and you should not worry about freeing the memory of the module yourself, but in some cases you may want to release the memory occupied by the module before the garbage collector steps in. In this case use the method [`delete`](../../06-api-reference/classes/ClassificationModule.md#forward) on the module object you will no longer use, and want to remove from the memory. Note that you cannot use [`forward`](../../06-api-reference/classes/ClassificationModule.md#forward) after [`delete`](../../06-api-reference/classes/ClassificationModule.md#forward) unless you load the module again.
+The module is a regular JavaScript object, and as such its lifespan will be managed by the garbage collector. In most cases this should be enough, and you should not worry about freeing the memory of the module yourself, but in some cases you may want to release the memory occupied by the module before the garbage collector steps in. In this case use the method [`delete`](../../06-api-reference/classes/ClassificationModule.md#delete) on the module object you will no longer use, and want to remove from the memory. Note that you cannot use [`forward`](../../06-api-reference/classes/ClassificationModule.md#forward) after [`delete`](../../06-api-reference/classes/ClassificationModule.md#delete) unless you load the module again.

packages/react-native-executorch/common/rnexecutorch/models/classification/Classification.cpp

@@ -1,18 +1,26 @@
 #include "Classification.h"
 
-#include <future>
-
 #include <rnexecutorch/Error.h>
 #include <rnexecutorch/ErrorCodes.h>
 #include <rnexecutorch/data_processing/ImageProcessing.h>
 #include <rnexecutorch/data_processing/Numerical.h>
-#include <rnexecutorch/models/classification/Constants.h>
 
 namespace rnexecutorch::models::classification {
 
 Classification::Classification(const std::string &modelSource,
+                               std::vector<float> normMean,
+                               std::vector<float> normStd,
+                               std::vector<std::string> labelNames,
                                std::shared_ptr<react::CallInvoker> callInvoker)
-    : VisionModel(modelSource, callInvoker) {
+    : VisionModel(modelSource, callInvoker),
+      labelNames_(std::move(labelNames)) {
+  if (normMean.size() == 3) {
+    normMean_ = cv::Scalar(normMean[0], normMean[1], normMean[2]);
+  }
+  if (normStd.size() == 3) {
+    normStd_ = cv::Scalar(normStd[0], normStd[1], normStd[2]);
+  }
+
   auto inputShapes = getAllInputShapes();
   if (inputShapes.size() == 0) {
     throw RnExecutorchError(RnExecutorchErrorCode::UnexpectedNumInputs,
@@ -37,7 +45,11 @@ Classification::runInference(cv::Mat image) {
   cv::Mat preprocessed = preprocess(image);
 
   auto inputTensor =
-      image_processing::getTensorFromMatrix(modelInputShape_, preprocessed);
+      (normMean_ && normStd_)
+          ? image_processing::getTensorFromMatrix(
+                modelInputShape_, preprocessed, *normMean_, *normStd_)
+          : image_processing::getTensorFromMatrix(modelInputShape_,
+                                                  preprocessed);
 
   auto forwardResult = BaseModel::forward(inputTensor);
   if (!forwardResult.ok()) {
@@ -78,13 +90,13 @@ Classification::postprocess(const Tensor &tensor) {
       static_cast<const float *>(tensor.const_data_ptr()), tensor.numel());
   std::vector<float> resultVec(resultData.begin(), resultData.end());
 
-  if (resultVec.size() != constants::kImagenet1kV1Labels.size()) {
+  if (resultVec.size() != labelNames_.size()) {
     char errorMessage[100];
     std::snprintf(
         errorMessage, sizeof(errorMessage),
         "Unexpected classification output size, was expecting: %zu classes "
         "but got: %zu classes",
-        constants::kImagenet1kV1Labels.size(), resultVec.size());
+        labelNames_.size(), resultVec.size());
     throw RnExecutorchError(RnExecutorchErrorCode::InvalidModelOutput,
                             errorMessage);
   }
@@ -93,7 +105,7 @@ Classification::postprocess(const Tensor &tensor) {
 
   std::unordered_map<std::string_view, float> probs;
   for (std::size_t cl = 0; cl < resultVec.size(); ++cl) {
-    probs[constants::kImagenet1kV1Labels[cl]] = resultVec[cl];
+    probs[labelNames_[cl]] = resultVec[cl];
   }
 
   return probs;

packages/react-native-executorch/common/rnexecutorch/models/classification/Classification.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <optional>
 #include <unordered_map>
 
 #include <executorch/extension/tensor/tensor_ptr.h>
@@ -16,7 +17,9 @@ using executorch::extension::TensorPtr;
 
 class Classification : public VisionModel {
 public:
-  Classification(const std::string &modelSource,
+  Classification(const std::string &modelSource, std::vector<float> normMean,
+                 std::vector<float> normStd,
+                 std::vector<std::string> labelNames,
                  std::shared_ptr<react::CallInvoker> callInvoker);
 
   [[nodiscard("Registered non-void function")]] std::unordered_map<
@@ -35,9 +38,15 @@ class Classification : public VisionModel {
   std::unordered_map<std::string_view, float> runInference(cv::Mat image);
 
   std::unordered_map<std::string_view, float> postprocess(const Tensor &tensor);
+
+  std::vector<std::string> labelNames_;
+  std::optional<cv::Scalar> normMean_;
+  std::optional<cv::Scalar> normStd_;
 };
 } // namespace models::classification
 
 REGISTER_CONSTRUCTOR(models::classification::Classification, std::string,
+                     std::vector<float>, std::vector<float>,
+                     std::vector<std::string>,
                      std::shared_ptr<react::CallInvoker>);
-} // namespace rnexecutorch
+} // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/tests/integration/ClassificationTest.cpp

@@ -15,6 +15,18 @@ constexpr auto kValidClassificationModelPath = "efficientnet_v2_s_xnnpack.pte";
 constexpr auto kValidTestImagePath =
     "file:///data/local/tmp/rnexecutorch_tests/test_image.jpg";
 
+static std::vector<float> kImagenetNormMean = {0.485f, 0.456f, 0.406f};
+static std::vector<float> kImagenetNormStd = {0.229f, 0.224f, 0.225f};
+
+static std::vector<std::string> getImagenetLabelNames() {
+  std::vector<std::string> names;
+  names.reserve(constants::kImagenet1kV1Labels.size());
+  for (const auto &label : constants::kImagenet1kV1Labels) {
+    names.emplace_back(label);
+  }
+  return names;
+}
+
 // ============================================================================
 // Common tests via typed test suite
 // ============================================================================
@@ -23,11 +35,12 @@ template <> struct ModelTraits<Classification> {
   using ModelType = Classification;
 
   static ModelType createValid() {
-    return ModelType(kValidClassificationModelPath, nullptr);
+    return ModelType(kValidClassificationModelPath, kImagenetNormMean,
+                     kImagenetNormStd, getImagenetLabelNames(), nullptr);
   }
 
   static ModelType createInvalid() {
-    return ModelType("nonexistent.pte", nullptr);
+    return ModelType("nonexistent.pte", {}, {}, {}, nullptr);
   }
 
   static void callGenerate(ModelType &model) {
@@ -46,37 +59,43 @@ INSTANTIATE_TYPED_TEST_SUITE_P(Classification, VisionModelTest,
 // Model-specific tests
 // ============================================================================
 TEST(ClassificationGenerateTests, InvalidImagePathThrows) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   EXPECT_THROW((void)model.generateFromString("nonexistent_image.jpg"),
                RnExecutorchError);
 }
 
 TEST(ClassificationGenerateTests, EmptyImagePathThrows) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   EXPECT_THROW((void)model.generateFromString(""), RnExecutorchError);
 }
 
 TEST(ClassificationGenerateTests, MalformedURIThrows) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   EXPECT_THROW((void)model.generateFromString("not_a_valid_uri://bad"),
                RnExecutorchError);
 }
 
 TEST(ClassificationGenerateTests, ValidImageReturnsResults) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto results = model.generateFromString(kValidTestImagePath);
   EXPECT_FALSE(results.empty());
 }
 
 TEST(ClassificationGenerateTests, ResultsHaveCorrectSize) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto results = model.generateFromString(kValidTestImagePath);
   auto expectedNumClasses = constants::kImagenet1kV1Labels.size();
   EXPECT_EQ(results.size(), expectedNumClasses);
 }
 
 TEST(ClassificationGenerateTests, ResultsContainValidProbabilities) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto results = model.generateFromString(kValidTestImagePath);
 
   float sum = 0.0f;
@@ -89,7 +108,8 @@ TEST(ClassificationGenerateTests, ResultsContainValidProbabilities) {
 }
 
 TEST(ClassificationGenerateTests, TopPredictionHasReasonableConfidence) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto results = model.generateFromString(kValidTestImagePath);
 
   float maxProb = 0.0f;
@@ -101,22 +121,32 @@ TEST(ClassificationGenerateTests, TopPredictionHasReasonableConfidence) {
   EXPECT_GT(maxProb, 0.0f);
 }
 
+TEST(ClassificationGenerateTests, WrongLabelCountThrows) {
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, {"A", "B", "C"}, nullptr);
+  EXPECT_THROW((void)model.generateFromString(kValidTestImagePath),
+               RnExecutorchError);
+}
+
 TEST(ClassificationInheritedTests, GetInputShapeWorks) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto shape = model.getInputShape("forward", 0);
   EXPECT_EQ(shape.size(), 4);
   EXPECT_EQ(shape[0], 1);
   EXPECT_EQ(shape[1], 3);
 }
 
 TEST(ClassificationInheritedTests, GetAllInputShapesWorks) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto shapes = model.getAllInputShapes("forward");
   EXPECT_FALSE(shapes.empty());
 }
 
 TEST(ClassificationInheritedTests, GetMethodMetaWorks) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   auto result = model.getMethodMeta("forward");
   EXPECT_TRUE(result.ok());
 }
@@ -125,7 +155,8 @@ TEST(ClassificationInheritedTests, GetMethodMetaWorks) {
 // generateFromPixels smoke test
 // ============================================================================
 TEST(ClassificationPixelTests, ValidPixelsReturnsResults) {
-  Classification model(kValidClassificationModelPath, nullptr);
+  Classification model(kValidClassificationModelPath, kImagenetNormMean,
+                       kImagenetNormStd, getImagenetLabelNames(), nullptr);
   std::vector<uint8_t> buf(64 * 64 * 3, 128);
   JSTensorViewIn view{
       buf.data(), {64, 64, 3}, executorch::aten::ScalarType::Byte};