70 lines
4.2 KiB
TypeScript
70 lines
4.2 KiB
TypeScript
import cv from "@yizhi/cv";
|
|
import { convertImage } from "../common/processors";
|
|
import { FaceRecognition, FaceRecognitionPredictOption } from "./common";
|
|
|
|
const MODEL_URL_CONFIG_ARC_FACE = {
|
|
INSIGHTFACE_ARCFACE_R100_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r100.onnx`,
|
|
INSIGHTFACE_ARCFACE_R50_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r50.onnx`,
|
|
INSIGHTFACE_ARCFACE_R34_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r34.onnx`,
|
|
INSIGHTFACE_ARCFACE_R18_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r18.onnx`,
|
|
INSIGHTFACE_ARCFACE_R50_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r50.mnn`,
|
|
INSIGHTFACE_ARCFACE_R34_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r34.mnn`,
|
|
INSIGHTFACE_ARCFACE_R18_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r18.mnn`,
|
|
};
|
|
const MODEL_URL_CONFIG_COS_FACE = {
|
|
INSIGHTFACE_COSFACE_R100_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r100.onnx`,
|
|
INSIGHTFACE_COSFACE_R50_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r50.onnx`,
|
|
INSIGHTFACE_COSFACE_R34_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r34.onnx`,
|
|
INSIGHTFACE_COSFACE_R18_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r18.onnx`,
|
|
INSIGHTFACE_COSFACE_R50_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r50.mnn`,
|
|
INSIGHTFACE_COSFACE_R34_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r34.mnn`,
|
|
INSIGHTFACE_COSFACE_R18_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r18.mnn`,
|
|
};
|
|
const MODEL_URL_CONFIG_PARTIAL_FC = {
|
|
INSIGHTFACE_PARTIALFC_R100_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/partial_fc_glint360k_r100.onnx`,
|
|
INSIGHTFACE_PARTIALFC_R50_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/partial_fc_glint360k_r50.onnx`,
|
|
INSIGHTFACE_PARTIALFC_R50_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/partial_fc_glint360k_r50.mnn`,
|
|
};
|
|
|
|
|
|
export class Insightface extends FaceRecognition {
|
|
|
|
public async doPredict(image: cv.Mat, option?: FaceRecognitionPredictOption): Promise<number[]> {
|
|
const input = this.input;
|
|
const output = this.output;
|
|
if (option?.crop) image = cv.crop(image, option.crop);
|
|
image = cv.resize(image, input.shape[3], input.shape[2]);
|
|
const nchwImageData = convertImage(image.data, { sourceImageFormat: "bgr", targetColorFormat: "bgr", targetShapeFormat: "nchw", targetNormalize: { mean: [127.5], std: [127.5] } });
|
|
|
|
const embedding = await this.session.run({
|
|
[input.name]: {
|
|
type: "float32",
|
|
data: nchwImageData,
|
|
shape: [1, 3, input.shape[2], input.shape[3]],
|
|
}
|
|
}).then(res => res[output.name]);
|
|
|
|
return new Array(...embedding.data);
|
|
}
|
|
|
|
}
|
|
|
|
export class ArcFace extends Insightface {
|
|
public static async load(type?: keyof typeof MODEL_URL_CONFIG_ARC_FACE) {
|
|
return this.cacheModel(MODEL_URL_CONFIG_ARC_FACE[type ?? "INSIGHTFACE_ARCFACE_R100_ONNX"], { createModel: true }).then(r => r.model);
|
|
}
|
|
}
|
|
|
|
export class CosFace extends Insightface {
|
|
public static async load(type?: keyof typeof MODEL_URL_CONFIG_COS_FACE) {
|
|
return this.cacheModel(MODEL_URL_CONFIG_COS_FACE[type ?? "INSIGHTFACE_COSFACE_R100_ONNX"], { createModel: true }).then(r => r.model);
|
|
}
|
|
}
|
|
|
|
export class PartialFC extends Insightface {
|
|
public static async load(type?: keyof typeof MODEL_URL_CONFIG_PARTIAL_FC) {
|
|
return this.cacheModel(MODEL_URL_CONFIG_PARTIAL_FC[type ?? "INSIGHTFACE_PARTIALFC_R100_ONNX"], { createModel: true }).then(r => r.model);
|
|
}
|
|
}
|
|
|