230 lines
8.0 KiB
TypeScript
230 lines
8.0 KiB
TypeScript
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";
|
|
|
|
export interface LocalUiBox {
|
|
x: number;
|
|
y: number;
|
|
width: number;
|
|
height: number;
|
|
category: string;
|
|
confidence: number;
|
|
}
|
|
|
|
interface DetectRequest {
|
|
id: number;
|
|
image: Buffer;
|
|
width: number;
|
|
height: number;
|
|
scoreThreshold?: number;
|
|
iouThreshold?: number;
|
|
}
|
|
|
|
interface DetectResponse {
|
|
id: number;
|
|
boxes?: LocalUiBox[];
|
|
error?: string;
|
|
elapsedMs?: number;
|
|
}
|
|
|
|
type ModelOutput = { output0: Tensor };
|
|
type TensorLike = { dims: number[]; data: Float32Array };
|
|
|
|
const inputName = "images";
|
|
const outputName = "output0";
|
|
const inputSize = 640;
|
|
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> {
|
|
if (!sessionPromise) {
|
|
const modelPath = String((workerData as { modelPath?: string }).modelPath || "");
|
|
sessionPromise = InferenceSession.create(modelPath, {
|
|
executionProviders: ["cpu"],
|
|
graphOptimizationLevel: "all",
|
|
executionMode: "sequential"
|
|
});
|
|
}
|
|
return sessionPromise;
|
|
}
|
|
|
|
function iou(a: LocalUiBox, b: LocalUiBox): number {
|
|
const x1 = Math.max(a.x, b.x);
|
|
const y1 = Math.max(a.y, b.y);
|
|
const x2 = Math.min(a.x + a.width, b.x + b.width);
|
|
const y2 = Math.min(a.y + a.height, b.y + b.height);
|
|
const intersection = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
|
|
if (intersection <= 0) return 0;
|
|
return intersection / (a.width * a.height + b.width * b.height - intersection);
|
|
}
|
|
|
|
function nms(boxes: LocalUiBox[], iouThreshold: number): LocalUiBox[] {
|
|
const sorted = [...boxes].sort((a, b) => b.confidence - a.confidence);
|
|
const kept: LocalUiBox[] = [];
|
|
for (const box of sorted) {
|
|
if (kept.every(item => iou(item, box) <= iouThreshold)) kept.push(box);
|
|
}
|
|
return kept;
|
|
}
|
|
|
|
function tensorAccessor(output: TensorLike, anchorCount: number, channelCount: number, channelFirst: boolean) {
|
|
return (channel: number, anchor: number): number => channelFirst
|
|
? output.data[channel * anchorCount + anchor]
|
|
: 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,
|
|
sourceHeight: number,
|
|
scoreThreshold: number,
|
|
iouThreshold: number
|
|
): LocalUiBox[] {
|
|
// YOLO/OmniParser exports are commonly [1,C,N] or [1,N,C]. C=4+classCount.
|
|
if (output.dims.length < 2) throw new Error(`不支持的模型输出维度：${output.dims.join("x")}`);
|
|
const [, first, second] = output.dims;
|
|
const channelFirst = first <= second;
|
|
const channelCount = channelFirst ? first : second;
|
|
const anchorCount = channelFirst ? second : first;
|
|
if (channelCount < 5) throw new Error(`模型输出通道数异常：${channelCount}`);
|
|
const at = tensorAccessor(output, anchorCount, channelCount, channelFirst);
|
|
const classCount = channelCount - 4;
|
|
|
|
const padded = Math.max(sourceWidth, sourceHeight);
|
|
const scale = inputSize / padded;
|
|
const drawnWidth = Math.round(sourceWidth * scale);
|
|
const drawnHeight = Math.round(sourceHeight * scale);
|
|
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++) {
|
|
let confidence = 0;
|
|
let classIndex = 0;
|
|
if (classCount === 1) {
|
|
confidence = at(4, index);
|
|
} else {
|
|
for (let candidate = 0; candidate < classCount; candidate++) {
|
|
const score = at(4 + candidate, index);
|
|
if (score > confidence) {
|
|
confidence = score;
|
|
classIndex = candidate;
|
|
}
|
|
}
|
|
}
|
|
if (!Number.isFinite(confidence) || confidence < scoreThreshold) continue;
|
|
const cx = at(0, index);
|
|
const cy = at(1, index);
|
|
const width = at(2, index);
|
|
const height = at(3, index);
|
|
let x = (cx - width / 2 - offsetX) / scale;
|
|
let y = (cy - height / 2 - offsetY) / scale;
|
|
let boxWidth = width / scale;
|
|
let boxHeight = height / scale;
|
|
x = Math.max(0, Math.min(sourceWidth - 1, x));
|
|
y = Math.max(0, Math.min(sourceHeight - 1, y));
|
|
boxWidth = Math.max(2, Math.min(sourceWidth - x, boxWidth));
|
|
boxHeight = Math.max(2, Math.min(sourceHeight - y, boxHeight));
|
|
if (boxWidth < minSide || boxHeight < minSide) continue;
|
|
const areaRatio = (boxWidth * boxHeight) / (sourceWidth * sourceHeight);
|
|
if (areaRatio > 0.92) continue;
|
|
raw.push({
|
|
x: Math.round(x),
|
|
y: Math.round(y),
|
|
width: Math.round(boxWidth),
|
|
height: Math.round(boxHeight),
|
|
category: labels?.[classIndex]
|
|
?? (classCount <= 1 ? "控件" : `类别${classIndex + 1}`),
|
|
confidence: Math.round(Math.max(0, Math.min(1, confidence)) * 1000) / 1000
|
|
});
|
|
}
|
|
|
|
return nms(raw, iouThreshold)
|
|
.sort((a, b) => b.confidence - a.confidence || a.y - b.y || a.x - b.x)
|
|
.slice(0, maxBoxes)
|
|
.sort((a, b) => a.y - b.y || a.x - b.x);
|
|
}
|
|
|
|
async function detect(request: DetectRequest): Promise<LocalUiBox[]> {
|
|
const model = await loadSession();
|
|
const source = await loadImage(request.image);
|
|
const width = request.width || source.width;
|
|
const height = request.height || source.height;
|
|
|
|
const canvas = createCanvas(inputSize, inputSize);
|
|
const context = canvas.getContext("2d");
|
|
context.fillStyle = "#727272";
|
|
context.fillRect(0, 0, inputSize, inputSize);
|
|
context.imageSmoothingEnabled = true;
|
|
context.imageSmoothingQuality = "high";
|
|
const scale = Math.min(inputSize / width, inputSize / height);
|
|
context.drawImage(
|
|
source,
|
|
(inputSize - width * scale) / 2,
|
|
(inputSize - height * scale) / 2,
|
|
width * scale,
|
|
height * scale
|
|
);
|
|
const pixels = context.getImageData(0, 0, inputSize, inputSize).data;
|
|
const tensorData = new Float32Array(3 * inputSize * inputSize);
|
|
const planeSize = inputSize * inputSize;
|
|
for (let index = 0; index < planeSize; index++) {
|
|
tensorData[index] = pixels[index * 4] / 255;
|
|
tensorData[planeSize + index] = pixels[index * 4 + 1] / 255;
|
|
tensorData[planeSize * 2 + index] = pixels[index * 4 + 2] / 255;
|
|
}
|
|
|
|
const results = await model.run({ [inputName]: new Tensor("float32", tensorData, [1, 3, inputSize, inputSize]) });
|
|
const output = results[outputName] as unknown as TensorLike;
|
|
return extractBoxes(
|
|
output,
|
|
width,
|
|
height,
|
|
request.scoreThreshold ?? defaultScoreThreshold,
|
|
request.iouThreshold ?? defaultIouThreshold
|
|
);
|
|
}
|
|
|
|
parentPort?.on("message", async (request: DetectRequest) => {
|
|
const started = Date.now();
|
|
try {
|
|
const boxes = await detect(request);
|
|
const response: DetectResponse = { id: request.id, boxes, elapsedMs: Date.now() - started };
|
|
parentPort?.postMessage(response);
|
|
} catch (error) {
|
|
const response: DetectResponse = {
|
|
id: request.id,
|
|
error: error instanceof Error ? error.message : String(error),
|
|
elapsedMs: Date.now() - started
|
|
};
|
|
parentPort?.postMessage(response);
|
|
}
|
|
});
|