chore: review suggestions

chmjkb · chmjkb · commit e2531bcc5551 · 2026-04-30T12:50:49.000+02:00
diff --git a/apps/computer-vision/app/pose_estimation/index.tsx b/apps/computer-vision/app/pose_estimation/index.tsx
@@ -155,6 +155,11 @@ export default function PoseEstimationScreen() {
                     (layout.width - imageDimensions.width * scaleX) / 2;
                   const offsetY =
                     (layout.height - imageDimensions.height * scaleY) / 2;
+                  const isInBounds = (kp: { x: number; y: number }) =>
+                    kp.x >= 0 &&
+                    kp.y >= 0 &&
+                    kp.x <= imageDimensions.width &&
+                    kp.y <= imageDimensions.height;
                   return (
                     <Svg style={StyleSheet.absoluteFill}>
                       {results.map((personKeypoints, personIdx) => {
@@ -167,6 +172,8 @@ export default function PoseEstimationScreen() {
                                 const kp1 = personKeypoints[from];
                                 const kp2 = personKeypoints[to];
                                 if (!kp1 || !kp2) return null;
+                                if (!isInBounds(kp1) || !isInBounds(kp2))
+                                  return null;
                                 return (
                                   <Line
                                     key={`person-${personIdx}-line-${lineIdx}`}
@@ -180,17 +187,17 @@ export default function PoseEstimationScreen() {
                                 );
                               }
                             )}
-                            {Object.entries(personKeypoints).map(
-                              ([name, kp]) => (
+                            {Object.entries(personKeypoints)
+                              .filter(([, kp]) => isInBounds(kp))
+                              .map(([name, kp]) => (
                                 <Circle
                                   key={`person-${personIdx}-kp-${name}`}
                                   cx={kp.x * scaleX + offsetX}
                                   cy={kp.y * scaleY + offsetY}
                                   r="4"
                                   fill="red"
                                 />
-                              )
-                            )}
+                              ))}
                           </React.Fragment>
                         );
                       })}
diff --git a/apps/computer-vision/components/vision_camera/tasks/PoseEstimationTask.tsx b/apps/computer-vision/components/vision_camera/tasks/PoseEstimationTask.tsx
@@ -142,13 +142,16 @@ export default function PoseEstimationTask({
       <Svg style={StyleSheet.absoluteFill}>
         {detections.map((personKeypoints, personIdx) => {
           const color = PERSON_COLORS[personIdx % PERSON_COLORS.length];
+          const isVisible = (kp: { x: number; y: number }) =>
+            kp.x >= 0 && kp.y >= 0;
           return (
             <React.Fragment key={`person-${personIdx}`}>
               {/* Draw skeleton lines */}
               {COCO_SKELETON_CONNECTIONS.map(([from, to], lineIdx) => {
                 const kp1 = personKeypoints[from];
                 const kp2 = personKeypoints[to];
                 if (!kp1 || !kp2) return null;
+                if (!isVisible(kp1) || !isVisible(kp2)) return null;
                 const x1 = kp1.x * scale + offsetX;
                 const y1 = kp1.y * scale + offsetY;
                 const x2 = kp2.x * scale + offsetX;
@@ -166,19 +169,21 @@ export default function PoseEstimationTask({
                 );
               })}
               {/* Draw keypoints */}
-              {Object.entries(personKeypoints).map(([name, kp]) => {
-                const cx = kp.x * scale + offsetX;
-                const cy = kp.y * scale + offsetY;
-                return (
-                  <Circle
-                    key={`person-${personIdx}-kp-${name}`}
-                    cx={cx}
-                    cy={cy}
-                    r={5}
-                    fill="red"
-                  />
-                );
-              })}
+              {Object.entries(personKeypoints)
+                .filter(([, kp]) => isVisible(kp))
+                .map(([name, kp]) => {
+                  const cx = kp.x * scale + offsetX;
+                  const cy = kp.y * scale + offsetY;
+                  return (
+                    <Circle
+                      key={`person-${personIdx}-kp-${name}`}
+                      cx={cx}
+                      cy={cy}
+                      r={5}
+                      fill="red"
+                    />
+                  );
+                })}
             </React.Fragment>
           );
         })}
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/pose_estimation/PoseEstimation.cpp b/packages/react-native-executorch/common/rnexecutorch/models/pose_estimation/PoseEstimation.cpp
@@ -32,7 +32,8 @@ PoseEstimation::PoseEstimation(const std::string &modelSource,
 
 PoseDetections PoseEstimation::postprocess(const std::vector<EValue> &tensors,
                                            cv::Size originalSize,
-                                           double detectionThreshold) {
+                                           double detectionThreshold,
+                                           double keypointThreshold) {
   // Output tensors (batch dim squeezed):
   //   0: boxes     (Q, 4)    - xyxy bbox in model input pixel space
   //   1: scores    (Q,)      - person confidence [0, 1]
@@ -75,6 +76,11 @@ PoseDetections PoseEstimation::postprocess(const std::vector<EValue> &tensors,
     const float *detectionKps = kpData + i * numKeypoints * 3;
 
     for (size_t k = 0; k < numKeypoints; ++k) {
+      float visibility = detectionKps[k * 3 + 2];
+      if (visibility < keypointThreshold) {
+        keypoints.emplace_back(-1, -1);
+        continue;
+      }
       float x = detectionKps[k * 3];
       float y = detectionKps[k * 3 + 1];
 
@@ -92,7 +98,7 @@ PoseDetections PoseEstimation::postprocess(const std::vector<EValue> &tensors,
 
 PoseDetections PoseEstimation::runInference(cv::Mat image,
                                             double detectionThreshold,
-                                            double iouThreshold,
+                                            double keypointThreshold,
                                             const std::string &methodName) {
 
   log(LOG_LEVEL::Debug, "Running inference with model name: " + methodName);
@@ -101,9 +107,9 @@ PoseDetections PoseEstimation::runInference(cv::Mat image,
     throw RnExecutorchError(RnExecutorchErrorCode::InvalidUserInput,
                             "detectionThreshold must be in range [0, 1]");
   }
-  if (iouThreshold < 0.0 || iouThreshold > 1.0) {
+  if (keypointThreshold < 0.0 || keypointThreshold > 1.0) {
     throw RnExecutorchError(RnExecutorchErrorCode::InvalidUserInput,
-                            "iouThreshold must be in range [0, 1]");
+                            "keypointThreshold must be in range [0, 1]");
   }
 
   std::scoped_lock lock(inference_mutex_);
@@ -132,30 +138,31 @@ PoseDetections PoseEstimation::runInference(cv::Mat image,
                                 "Ensure the model input is correct.");
   }
 
-  return postprocess(executeResult.get(), originalSize, detectionThreshold);
+  return postprocess(executeResult.get(), originalSize, detectionThreshold,
+                     keypointThreshold);
 }
 
 PoseDetections PoseEstimation::generateFromString(std::string imageSource,
                                                   double detectionThreshold,
-                                                  double iouThreshold,
+                                                  double keypointThreshold,
                                                   std::string methodName) {
   cv::Mat imageBGR = image_processing::readImage(imageSource);
   cv::Mat imageRGB;
   cv::cvtColor(imageBGR, imageRGB, cv::COLOR_BGR2RGB);
-  return runInference(std::move(imageRGB), detectionThreshold, iouThreshold,
-                      methodName);
+  return runInference(std::move(imageRGB), detectionThreshold,
+                      keypointThreshold, methodName);
 }
 
 PoseDetections PoseEstimation::generateFromFrame(jsi::Runtime &runtime,
                                                  const jsi::Value &frameData,
                                                  double detectionThreshold,
-                                                 double iouThreshold,
+                                                 double keypointThreshold,
                                                  std::string methodName) {
   auto orient = ::rnexecutorch::utils::readFrameOrientation(runtime, frameData);
   cv::Mat frame = extractFromFrame(runtime, frameData);
   cv::Mat rotated = ::rnexecutorch::utils::rotateFrameForModel(frame, orient);
   auto detections =
-      runInference(rotated, detectionThreshold, iouThreshold, methodName);
+      runInference(rotated, detectionThreshold, keypointThreshold, methodName);
   for (auto &person : detections) {
     ::rnexecutorch::utils::inverseRotatePoints(person, orient, rotated.size());
   }
@@ -164,10 +171,10 @@ PoseDetections PoseEstimation::generateFromFrame(jsi::Runtime &runtime,
 
 PoseDetections PoseEstimation::generateFromPixels(JSTensorViewIn pixelData,
                                                   double detectionThreshold,
-                                                  double iouThreshold,
+                                                  double keypointThreshold,
                                                   std::string methodName) {
   cv::Mat image = extractFromPixels(pixelData);
-  return runInference(image, detectionThreshold, iouThreshold, methodName);
+  return runInference(image, detectionThreshold, keypointThreshold, methodName);
 }
 
 } // namespace rnexecutorch::models::pose_estimation
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/pose_estimation/PoseEstimation.h b/packages/react-native-executorch/common/rnexecutorch/models/pose_estimation/PoseEstimation.h
@@ -17,27 +17,28 @@ class PoseEstimation : public VisionModel {
 
   [[nodiscard("Registered non-void function")]] PoseDetections
   generateFromString(std::string imageSource, double detectionThreshold,
-                     double iouThreshold, std::string methodName);
+                     double keypointThreshold, std::string methodName);
   [[nodiscard("Registered non-void function")]] PoseDetections
   generateFromFrame(jsi::Runtime &runtime, const jsi::Value &frameData,
-                    double detectionThreshold, double iouThreshold,
+                    double detectionThreshold, double keypointThreshold,
                     std::string methodName);
   [[nodiscard("Registered non-void function")]] PoseDetections
   generateFromPixels(JSTensorViewIn pixelData, double detectionThreshold,
-                     double iouThreshold, std::string methodName);
+                     double keypointThreshold, std::string methodName);
 
 private:
   std::optional<cv::Scalar> normMean_;
   std::optional<cv::Scalar> normStd_;
 
   [[nodiscard("Registered non-void function")]]
   PoseDetections runInference(cv::Mat image, double detectionThreshold,
-                              double iouThreshold,
+                              double keypointThreshold,
                               const std::string &modelName);
 
   [[nodiscard("Registered non-void function")]]
   PoseDetections postprocess(const std::vector<EValue> &evl,
-                             cv::Size originalSize, double detectionThreshold);
+                             cv::Size originalSize, double detectionThreshold,
+                             double keypointThreshold);
 };
 
 } // namespace models::pose_estimation
diff --git a/packages/react-native-executorch/common/rnexecutorch/tests/integration/PoseEstimationTest.cpp b/packages/react-native-executorch/common/rnexecutorch/tests/integration/PoseEstimationTest.cpp
@@ -82,14 +82,14 @@ TEST(PoseEstimationGenerateTests, DetectionThresholdAboveOneThrows) {
                RnExecutorchError);
 }
 
-TEST(PoseEstimationGenerateTests, NegativeIouThresholdThrows) {
+TEST(PoseEstimationGenerateTests, NegativeKeypointThresholdThrows) {
   PoseEstimation model(kValidPoseModelPath, {}, {}, nullptr);
   EXPECT_THROW((void)model.generateFromString(kValidTestImagePath, 0.5, -0.1,
                                               kMethodName),
                RnExecutorchError);
 }
 
-TEST(PoseEstimationGenerateTests, IouThresholdAboveOneThrows) {
+TEST(PoseEstimationGenerateTests, KeypointThresholdAboveOneThrows) {
   PoseEstimation model(kValidPoseModelPath, {}, {}, nullptr);
   EXPECT_THROW((void)model.generateFromString(kValidTestImagePath, 0.5, 1.1,
                                               kMethodName),
diff --git a/packages/react-native-executorch/src/modules/computer_vision/PoseEstimationModule.ts b/packages/react-native-executorch/src/modules/computer_vision/PoseEstimationModule.ts
@@ -22,7 +22,7 @@ const YOLO_POSE_CONFIG = {
   availableInputSizes: [384, 512, 640] as const,
   defaultInputSize: 384,
   defaultDetectionThreshold: 0.5,
-  defaultIouThreshold: 0.5,
+  defaultKeypointThreshold: 0.5,
 } satisfies PoseEstimationConfig<typeof CocoKeypoint>;
 
 const ModelConfigs = {
@@ -48,6 +48,26 @@ type ModelNameOf<C extends PoseEstimationModelSources> = C['modelName'];
 type ResolveKeypoints<T extends PoseEstimationModelName | KeypointEnum> =
   ResolveConfigOrType<T, ModelConfigsType, 'keypointMap'>;
 
+function mapPersonKeypoints<K extends KeypointEnum>(
+  raw: Keypoint[][],
+  entries: [string, number][],
+  maxIndex: number
+): PersonKeypoints<K>[] {
+  'worklet';
+  if (raw.length > 0 && raw[0]!.length <= maxIndex) {
+    throw new Error(
+      `Keypoint map references index ${maxIndex} but model returned ${raw[0]!.length} keypoints per person — keypointMap is incompatible with this model.`
+    );
+  }
+  const out: PersonKeypoints<K>[] = [];
+  for (const person of raw) {
+    const named: Record<string, Keypoint> = {};
+    for (const [name, idx] of entries) named[name] = person[idx]!;
+    out.push(named as PersonKeypoints<K>);
+  }
+  return out;
+}
+
 /**
  * Pose estimation module for detecting human body keypoints.
  * @typeParam T - Either a built-in model name (e.g. `'yolo26n-pose'`)
@@ -59,6 +79,7 @@ export class PoseEstimationModule<
 > extends VisionModule<PoseDetections<ResolveKeypoints<T>>> {
   private readonly keypointMap: ResolveKeypoints<T>;
   private readonly modelConfig: PoseEstimationConfig<KeypointEnum>;
+  private readonly maxKeypointIndex: number;
 
   private constructor(
     keypointMap: ResolveKeypoints<T>,
@@ -69,6 +90,7 @@ export class PoseEstimationModule<
     this.keypointMap = keypointMap;
     this.modelConfig = modelConfig;
     this.nativeModule = nativeModule;
+    this.maxKeypointIndex = Math.max(...Object.values(keypointMap));
   }
 
   /**
@@ -169,14 +191,12 @@ export class PoseEstimationModule<
     const nativeGenerateFromFrame = this.nativeModule.generateFromFrame;
     const defaultDetectionThreshold =
       this.modelConfig.defaultDetectionThreshold ?? 0.5;
-    const defaultIouThreshold = this.modelConfig.defaultIouThreshold ?? 0.5;
+    const defaultKeypointThreshold =
+      this.modelConfig.defaultKeypointThreshold ?? 0.5;
     const defaultInputSize = this.modelConfig.defaultInputSize;
     const availableInputSizes = this.modelConfig.availableInputSizes;
-    const keypointEntries = Object.entries(this.keypointMap) as [
-      string,
-      number,
-    ][];
-
+    const keypointEntries = Object.entries(this.keypointMap);
+    const maxKeypointIndex = this.maxKeypointIndex;
     return (
       frame: Frame,
       isFrontCamera: boolean,
@@ -186,7 +206,8 @@ export class PoseEstimationModule<
 
       const detectionThreshold =
         options?.detectionThreshold ?? defaultDetectionThreshold;
-      const iouThreshold = options?.iouThreshold ?? defaultIouThreshold;
+      const keypointThreshold =
+        options?.keypointThreshold ?? defaultKeypointThreshold;
       const inputSize = options?.inputSize ?? defaultInputSize;
 
       // Validate inputSize
@@ -214,16 +235,14 @@ export class PoseEstimationModule<
         const raw: Keypoint[][] = nativeGenerateFromFrame(
           frameData,
           detectionThreshold,
-          iouThreshold,
+          keypointThreshold,
           methodName
         );
-        const out: PersonKeypoints<ResolveKeypoints<T>>[] = [];
-        for (const person of raw) {
-          const named: Record<string, Keypoint> = {};
-          for (const [name, idx] of keypointEntries) named[name] = person[idx]!;
-          out.push(named as PersonKeypoints<ResolveKeypoints<T>>);
-        }
-        return out;
+        return mapPersonKeypoints<ResolveKeypoints<T>>(
+          raw,
+          keypointEntries,
+          maxKeypointIndex
+        );
       } finally {
         if (nativeBuffer?.release) {
           nativeBuffer.release();
@@ -253,8 +272,10 @@ export class PoseEstimationModule<
       options?.detectionThreshold ??
       this.modelConfig.defaultDetectionThreshold ??
       0.5;
-    const iouThreshold =
-      options?.iouThreshold ?? this.modelConfig.defaultIouThreshold ?? 0.5;
+    const keypointThreshold =
+      options?.keypointThreshold ??
+      this.modelConfig.defaultKeypointThreshold ??
+      0.5;
     const inputSize = options?.inputSize ?? this.modelConfig.defaultInputSize;
 
     // Validate inputSize against availableInputSizes
@@ -277,21 +298,21 @@ export class PoseEstimationModule<
         ? await this.nativeModule.generateFromString(
             input,
             detectionThreshold,
-            iouThreshold,
+            keypointThreshold,
             methodName
           )
         : await this.nativeModule.generateFromPixels(
             input,
             detectionThreshold,
-            iouThreshold,
+            keypointThreshold,
             methodName
           );
 
-    const entries = Object.entries(this.keypointMap) as [string, number][];
-    return raw.map((person) => {
-      const named: Record<string, Keypoint> = {};
-      for (const [name, idx] of entries) named[name] = person[idx]!;
-      return named as PersonKeypoints<ResolveKeypoints<T>>;
-    });
+    const entries = Object.entries(this.keypointMap);
+    return mapPersonKeypoints<ResolveKeypoints<T>>(
+      raw,
+      entries,
+      this.maxKeypointIndex
+    );
   }
 }
diff --git a/packages/react-native-executorch/src/types/poseEstimation.ts b/packages/react-native-executorch/src/types/poseEstimation.ts
