feat: enhance AI pane picking with OCR and cancellable hotkeys

src/main/local-vision/ui-detect-worker.ts

@@ -1,3 +1,4 @@
+import fs from "node:fs";
 import { parentPort, workerData } from "node:worker_threads";
 import { InferenceSession, Tensor } from "onnxruntime-node";
 import { createCanvas, loadImage } from "@napi-rs/canvas";
@@ -37,6 +38,9 @@ const defaultScoreThreshold = 0.25;
 const defaultIouThreshold = 0.42;
 const maxBoxes = 160;
 
+const startupModelPath = String((workerData as { modelPath?: string }).modelPath || "");
+const modelLabels = readModelLabels(startupModelPath);
+
 let sessionPromise: Promise<InferenceSession> | null = null;
 
 function loadSession(): Promise<InferenceSession> {
@@ -76,6 +80,24 @@ function tensorAccessor(output: TensorLike, anchorCount: number, channelCount: n
     : output.data[anchor * channelCount + channel];
 }
 
+// YOLO 导出的 ONNX 通常不含标签元数据。约定同名 *.labels.json 保存类别：
+// 1) ["button", "input"] 2) { "labels": ["button", "input"] }
+function readModelLabels(modelPath: string): string[] | undefined {
+  const sidecarPath = modelPath.replace(/\.onnx$/i, ".labels.json");
+  try {
+    const parsed: unknown = JSON.parse(fs.readFileSync(sidecarPath, "utf8"));
+    const values = Array.isArray(parsed) ? parsed : (parsed as { labels?: unknown })?.labels;
+    if (!Array.isArray(values)) return undefined;
+    const labels = values
+      .filter((item): item is string | number => typeof item === "string" || typeof item === "number")
+      .map(item => String(item))
+      .filter(Boolean);
+    return labels.length ? labels : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
 function extractBoxes(
   output: TensorLike,
   sourceWidth: number,
@@ -100,6 +122,7 @@ function extractBoxes(
   const offsetX = (inputSize - drawnWidth) / 2;
   const offsetY = (inputSize - drawnHeight) / 2;
   const minSide = Math.max(8, Math.min(sourceWidth, sourceHeight) * 0.006);
+  const labels = modelLabels;
   const raw: LocalUiBox[] = [];
 
   for (let index = 0; index < anchorCount; index++) {
@@ -137,7 +160,8 @@ function extractBoxes(
       y: Math.round(y),
       width: Math.round(boxWidth),
       height: Math.round(boxHeight),
-      category: classCount <= 1 ? "控件" : `类别${classIndex + 1}`,
+      category: labels?.[classIndex]
+        ?? (classCount <= 1 ? "控件" : `类别${classIndex + 1}`),
       confidence: Math.round(Math.max(0, Math.min(1, confidence)) * 1000) / 1000
     });
   }