300 lines
14 KiB
TypeScript
300 lines
14 KiB
TypeScript
import fs from "node:fs/promises";
|
|
|
|
import * as ort from "onnxruntime-node";
|
|
|
|
import type { DecodedPortraitImage } from "./image-metrics.js";
|
|
import { cosineSimilarity } from "./image-metrics.js";
|
|
import { portraitModelDirectory, portraitModelStatus } from "./model-store.js";
|
|
|
|
export type PortraitFace = {
|
|
bbox: [number, number, number, number];
|
|
detScore: number;
|
|
keypoints: Array<[number, number]>;
|
|
areaRatio: number;
|
|
embedding?: Float32Array;
|
|
};
|
|
|
|
export type PortraitFaceAnalysis = {
|
|
faces: PortraitFace[];
|
|
selectedFace?: PortraitFace;
|
|
embedding?: Float32Array;
|
|
};
|
|
|
|
export class PortraitFaceEngineError extends Error {
|
|
constructor(readonly code: "portrait_model_missing" | "portrait_face_engine_failed", message: string) {
|
|
super(message);
|
|
this.name = "PortraitFaceEngineError";
|
|
}
|
|
}
|
|
|
|
type FaceEngineOptions = {
|
|
modelRoot: string;
|
|
detectorSize?: number;
|
|
threshold?: number;
|
|
};
|
|
|
|
type DetectorSession = { session: ort.InferenceSession; inputName: string };
|
|
type EmbeddingSession = { session: ort.InferenceSession; inputName: string };
|
|
|
|
export class PortraitFaceEngine {
|
|
private detector?: Promise<DetectorSession>;
|
|
private embedding?: Promise<EmbeddingSession>;
|
|
private readonly detectorSize: number;
|
|
private readonly threshold: number;
|
|
|
|
constructor(private readonly options: FaceEngineOptions) {
|
|
this.detectorSize = options.detectorSize ?? 640;
|
|
// Keep low-confidence detections available for crop/search and let the
|
|
// precheck mark them unreliable instead of discarding them before the
|
|
// report can explain why the result is inconclusive.
|
|
this.threshold = options.threshold ?? 0.4;
|
|
}
|
|
|
|
async status() {
|
|
return portraitModelStatus(this.options.modelRoot);
|
|
}
|
|
|
|
async analyze(image: DecodedPortraitImage): Promise<PortraitFaceAnalysis> {
|
|
const models = await this.status();
|
|
if (!models.ready) throw new PortraitFaceEngineError("portrait_model_missing", "本地肖像 ONNX 模型尚未安装");
|
|
try {
|
|
const faces = await this.detect(image);
|
|
for (const face of faces) face.embedding = await this.embed(image, face.keypoints);
|
|
const selectedFace = [...faces].sort((left, right) => right.detScore - left.detScore || right.areaRatio - left.areaRatio)[0];
|
|
return { faces, selectedFace, embedding: selectedFace?.embedding };
|
|
} catch (error) {
|
|
if (error instanceof PortraitFaceEngineError) throw error;
|
|
throw new PortraitFaceEngineError("portrait_face_engine_failed", "本地人脸模型执行失败");
|
|
}
|
|
}
|
|
|
|
async compare(left: DecodedPortraitImage, right: DecodedPortraitImage) {
|
|
const [a, b] = await Promise.all([this.analyze(left), this.analyze(right)]);
|
|
return { left: a, right: b, similarity: a.embedding && b.embedding ? cosineSimilarity(a.embedding, b.embedding) : undefined };
|
|
}
|
|
|
|
private async detect(image: DecodedPortraitImage) {
|
|
const session = await this.getDetectorSession();
|
|
const input = await detectorTensor(image, this.detectorSize);
|
|
const output = await session.session.run({ [session.inputName]: input.tensor });
|
|
const detections: PortraitFace[] = [];
|
|
const outputs = session.session.outputNames.map((name) => output[name]);
|
|
if (outputs.length < 9) throw new Error("portrait_detector_outputs_invalid");
|
|
for (let level = 0; level < 3; level += 1) {
|
|
const scores = numericTensorData(outputs[level]!);
|
|
const boxes = numericTensorData(outputs[level + 3]!);
|
|
const keypoints = numericTensorData(outputs[level + 6]!);
|
|
const anchors = scores.length;
|
|
const gridWidth = Math.round(Math.sqrt(anchors / 2));
|
|
const gridHeight = gridWidth;
|
|
if (gridWidth * gridHeight * 2 !== anchors || boxes.length !== anchors * 4 || keypoints.length !== anchors * 10) throw new Error("portrait_detector_shape_invalid");
|
|
const stride = [8, 16, 32][level]!;
|
|
for (let index = 0; index < anchors; index += 1) {
|
|
const score = normalizedScore(Number(scores[index]));
|
|
if (score < this.threshold) continue;
|
|
const cell = Math.floor(index / 2);
|
|
const x = cell % gridWidth;
|
|
const y = Math.floor(cell / gridWidth);
|
|
const centerX = x * stride;
|
|
const centerY = y * stride;
|
|
const offset = index * 4;
|
|
const left = centerX - Number(boxes[offset]) * stride;
|
|
const top = centerY - Number(boxes[offset + 1]) * stride;
|
|
const right = centerX + Number(boxes[offset + 2]) * stride;
|
|
const bottom = centerY + Number(boxes[offset + 3]) * stride;
|
|
const pointOffset = index * 10;
|
|
const points: Array<[number, number]> = [];
|
|
for (let point = 0; point < 5; point += 1) points.push([centerX + Number(keypoints[pointOffset + point * 2]) * stride, centerY + Number(keypoints[pointOffset + point * 2 + 1]) * stride]);
|
|
detections.push({
|
|
bbox: mapBoxToSource([left, top, right, bottom], image, input.scale, input.offsetX, input.offsetY),
|
|
detScore: score,
|
|
keypoints: points.map(([pointX, pointY]) => mapPointToSource(pointX, pointY, image, input.scale, input.offsetX, input.offsetY)),
|
|
areaRatio: 0,
|
|
});
|
|
}
|
|
}
|
|
const selected = nonMaximumSuppression(detections, 0.4);
|
|
return selected.map((face) => ({ ...face, areaRatio: areaRatio(face.bbox, image.width, image.height) }));
|
|
}
|
|
|
|
private async embed(image: DecodedPortraitImage, keypoints: Array<[number, number]>) {
|
|
const session = await this.getEmbeddingSession();
|
|
const aligned = alignFace(image, keypoints);
|
|
const data = new Float32Array(1 * 3 * 112 * 112);
|
|
for (let index = 0; index < 112 * 112; index += 1) {
|
|
const source = index * 3;
|
|
data[index] = (aligned[source]! - 127.5) / 127.5;
|
|
data[112 * 112 + index] = (aligned[source + 1]! - 127.5) / 127.5;
|
|
data[2 * 112 * 112 + index] = (aligned[source + 2]! - 127.5) / 127.5;
|
|
}
|
|
const output = await session.session.run({ [session.inputName]: new ort.Tensor("float32", data, [1, 3, 112, 112]) });
|
|
const embedding = new Float32Array(numericTensorData(output[session.session.outputNames[0]!]));
|
|
let norm = 0;
|
|
for (const value of embedding) norm += value * value;
|
|
norm = Math.sqrt(norm);
|
|
if (!norm) return embedding;
|
|
for (let index = 0; index < embedding.length; index += 1) embedding[index] = embedding[index]! / norm;
|
|
return embedding;
|
|
}
|
|
|
|
private getDetectorSession() {
|
|
this.detector ??= this.loadDetector();
|
|
return this.detector;
|
|
}
|
|
|
|
private getEmbeddingSession() {
|
|
this.embedding ??= this.loadEmbedding();
|
|
return this.embedding;
|
|
}
|
|
|
|
private async loadDetector() {
|
|
const filePath = `${portraitModelDirectory(this.options.modelRoot)}/det_10g.onnx`;
|
|
await fs.access(filePath);
|
|
const session = await ort.InferenceSession.create(filePath, { executionProviders: ["cpu"] });
|
|
return { session, inputName: session.inputNames[0]! };
|
|
}
|
|
|
|
private async loadEmbedding() {
|
|
const filePath = `${portraitModelDirectory(this.options.modelRoot)}/w600k_r50.onnx`;
|
|
await fs.access(filePath);
|
|
const session = await ort.InferenceSession.create(filePath, { executionProviders: ["cpu"] });
|
|
return { session, inputName: session.inputNames[0]! };
|
|
}
|
|
}
|
|
|
|
function normalizedScore(value: number) {
|
|
if (!Number.isFinite(value)) return 0;
|
|
return value >= 0 && value <= 1 ? value : 1 / (1 + Math.exp(-value));
|
|
}
|
|
|
|
function numericTensorData(value: unknown) {
|
|
if (!value || typeof value !== "object" || !("data" in value)) throw new Error("portrait_tensor_invalid");
|
|
const data = (value as { data: unknown }).data;
|
|
if (!data || typeof (data as ArrayLike<unknown>).length !== "number") throw new Error("portrait_tensor_invalid");
|
|
return data as ArrayLike<number>;
|
|
}
|
|
|
|
async function detectorTensor(image: DecodedPortraitImage, size: number) {
|
|
const source = await import("sharp");
|
|
const resized = await source.default(image.rgb, { raw: { width: image.width, height: image.height, channels: 3 } })
|
|
.resize(size, size, { fit: "contain", background: { r: 0, g: 0, b: 0 } })
|
|
.raw()
|
|
.toBuffer({ resolveWithObject: true });
|
|
const scale = Math.min(size / image.width, size / image.height);
|
|
const offsetX = (size - image.width * scale) / 2;
|
|
const offsetY = (size - image.height * scale) / 2;
|
|
const data = new Float32Array(1 * 3 * size * size);
|
|
for (let index = 0; index < size * size; index += 1) {
|
|
const sourceIndex = index * 3;
|
|
data[index] = (resized.data[sourceIndex]! - 127.5) / 128;
|
|
data[size * size + index] = (resized.data[sourceIndex + 1]! - 127.5) / 128;
|
|
data[2 * size * size + index] = (resized.data[sourceIndex + 2]! - 127.5) / 128;
|
|
}
|
|
return { tensor: new ort.Tensor("float32", data, [1, 3, size, size]), scale, offsetX, offsetY };
|
|
}
|
|
|
|
function mapBoxToSource(box: [number, number, number, number], image: DecodedPortraitImage, scale: number, offsetX: number, offsetY: number): [number, number, number, number] {
|
|
return [
|
|
clamp((box[0] - offsetX) / scale, 0, image.width),
|
|
clamp((box[1] - offsetY) / scale, 0, image.height),
|
|
clamp((box[2] - offsetX) / scale, 0, image.width),
|
|
clamp((box[3] - offsetY) / scale, 0, image.height),
|
|
];
|
|
}
|
|
|
|
function mapPointToSource(x: number, y: number, image: DecodedPortraitImage, scale: number, offsetX: number, offsetY: number): [number, number] {
|
|
return [clamp((x - offsetX) / scale, 0, image.width), clamp((y - offsetY) / scale, 0, image.height)];
|
|
}
|
|
|
|
function nonMaximumSuppression(faces: PortraitFace[], threshold: number) {
|
|
const sorted = [...faces].sort((left, right) => right.detScore - left.detScore);
|
|
const kept: PortraitFace[] = [];
|
|
while (sorted.length) {
|
|
const current = sorted.shift()!;
|
|
kept.push(current);
|
|
for (let index = sorted.length - 1; index >= 0; index -= 1) if (intersectionOverUnion(current.bbox, sorted[index]!.bbox) > threshold) sorted.splice(index, 1);
|
|
}
|
|
return kept;
|
|
}
|
|
|
|
function intersectionOverUnion(a: PortraitFace["bbox"], b: PortraitFace["bbox"]) {
|
|
const left = Math.max(a[0], b[0]);
|
|
const top = Math.max(a[1], b[1]);
|
|
const right = Math.min(a[2], b[2]);
|
|
const bottom = Math.min(a[3], b[3]);
|
|
const intersection = Math.max(0, right - left) * Math.max(0, bottom - top);
|
|
const areaA = Math.max(0, a[2] - a[0]) * Math.max(0, a[3] - a[1]);
|
|
const areaB = Math.max(0, b[2] - b[0]) * Math.max(0, b[3] - b[1]);
|
|
return intersection / Math.max(1, areaA + areaB - intersection);
|
|
}
|
|
|
|
function areaRatio(box: PortraitFace["bbox"], width: number, height: number) {
|
|
return Math.max(0, box[2] - box[0]) * Math.max(0, box[3] - box[1]) / Math.max(1, width * height);
|
|
}
|
|
|
|
function alignFace(image: DecodedPortraitImage, keypoints: Array<[number, number]>) {
|
|
const target: Array<[number, number]> = [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]];
|
|
const transform = similarityTransform(keypoints, target);
|
|
const aligned = new Uint8Array(112 * 112 * 3);
|
|
for (let y = 0; y < 112; y += 1) {
|
|
for (let x = 0; x < 112; x += 1) {
|
|
const sourcePoint = inverseTransform(transform, x, y);
|
|
const targetIndex = (y * 112 + x) * 3;
|
|
sampleRgb(image, sourcePoint[0], sourcePoint[1], aligned, targetIndex);
|
|
}
|
|
}
|
|
return aligned;
|
|
}
|
|
|
|
function similarityTransform(source: Array<[number, number]>, target: Array<[number, number]>) {
|
|
const sourceMean = meanPoint(source);
|
|
const targetMean = meanPoint(target);
|
|
let alpha = 0;
|
|
let beta = 0;
|
|
let denominator = 0;
|
|
for (let index = 0; index < source.length; index += 1) {
|
|
const sx = source[index]![0] - sourceMean[0];
|
|
const sy = source[index]![1] - sourceMean[1];
|
|
const tx = target[index]![0] - targetMean[0];
|
|
const ty = target[index]![1] - targetMean[1];
|
|
alpha += sx * tx + sy * ty;
|
|
beta += sx * ty - sy * tx;
|
|
denominator += sx * sx + sy * sy;
|
|
}
|
|
alpha /= denominator || 1;
|
|
beta /= denominator || 1;
|
|
return {
|
|
alpha,
|
|
beta,
|
|
tx: targetMean[0] - alpha * sourceMean[0] + beta * sourceMean[1],
|
|
ty: targetMean[1] - beta * sourceMean[0] - alpha * sourceMean[1],
|
|
};
|
|
}
|
|
|
|
function inverseTransform(transform: { alpha: number; beta: number; tx: number; ty: number }, x: number, y: number): [number, number] {
|
|
const dx = x - transform.tx;
|
|
const dy = y - transform.ty;
|
|
const denominator = transform.alpha * transform.alpha + transform.beta * transform.beta || 1;
|
|
return [
|
|
(transform.alpha * dx + transform.beta * dy) / denominator,
|
|
(-transform.beta * dx + transform.alpha * dy) / denominator,
|
|
];
|
|
}
|
|
|
|
function sampleRgb(image: DecodedPortraitImage, x: number, y: number, target: Uint8Array, offset: number) {
|
|
const x0 = clamp(Math.floor(x), 0, image.width - 1);
|
|
const y0 = clamp(Math.floor(y), 0, image.height - 1);
|
|
const index = (y0 * image.width + x0) * 3;
|
|
target[offset] = image.rgb[index]!;
|
|
target[offset + 1] = image.rgb[index + 1]!;
|
|
target[offset + 2] = image.rgb[index + 2]!;
|
|
}
|
|
|
|
function meanPoint(points: Array<[number, number]>): [number, number] {
|
|
return [points.reduce((sum, point) => sum + point[0], 0) / points.length, points.reduce((sum, point) => sum + point[1], 0) / points.length];
|
|
}
|
|
|
|
function clamp(value: number, min: number, max: number) {
|
|
return Math.min(max, Math.max(min, value));
|
|
}
|