import cv from "@yizhi/cv"; import { ImageCropOption, ImageSource, Model } from "../common/model"; import { convertImage } from "../common/processors"; import { FaceAlignmentResult, FacePoint } from "./common"; export interface PFLDPredictOption extends ImageCropOption { } class PFLDResult extends FaceAlignmentResult { public constructor(points: FacePoint[]) { super(points, { directionPointIndex: [36, 92], leftEyePointIndex: [33, 34, 35, 36, 37, 38, 39, 40, 41, 42], rightEyePointIndex: [87, 88, 89, 90, 91, 92, 93, 94, 95, 96], leftEyebrowPointIndex: [43, 44, 45, 46, 47, 48, 49, 50, 51], rightEyebrowPointIndex: [97, 98, 99, 100, 101, 102, 103, 104, 105], mouthPointIndex: [52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], nosePointIndex: [72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86], contourPointIndex: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32], }); } } const MODEL_URL_CONFIG = { PFLD_106_LITE_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/facealign/pfld-106-lite.onnx`, PFLD_106_V2_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/facealign/pfld-106-v2.onnx`, PFLD_106_V3_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/facealign/pfld-106-v3.onnx`, PFLD_106_LITE_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/facealign/pfld-106-lite.mnn`, PFLD_106_V2_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/facealign/pfld-106-v2.mnn`, PFLD_106_V3_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/facealign/pfld-106-v3.mnn`, }; export class PFLD extends Model { public static async load(type?: keyof typeof MODEL_URL_CONFIG) { return this.cacheModel(MODEL_URL_CONFIG[type ?? "PFLD_106_LITE_ONNX"], { createModel: true }).then(r => r.model); } public predict(image: ImageSource, option?: PFLDPredictOption) { return Model.resolveImage(image, im => this.doPredict(im, option)); } private async doPredict(image: cv.Mat, option?: PFLDPredictOption) { const input = this.input; if (option?.crop) image = cv.crop(image, option.crop); const ratioWidth = image.cols / input.shape[3]; const ratioHeight = image.rows / input.shape[2]; image = cv.resize(image, input.shape[3], input.shape[2]); const nchwImageData = convertImage(image.data, { sourceImageFormat: "bgr", targetColorFormat: "bgr", targetShapeFormat: "nchw", targetNormalize: { mean: [0], std: [255] } }) const pointsOutput = Object.entries(this.outputs).filter(([_, out]) => out.shape.length == 2)[0][1]; const res = await this.session.run({ [input.name]: { type: "float32", data: nchwImageData, shape: [1, 3, input.shape[2], input.shape[3]], } }); const pointsBuffer = res[pointsOutput.name].data; const points: FacePoint[] = []; for (let i = 0; i < pointsBuffer.length; i += 2) { const x = pointsBuffer[i] * image.cols * ratioWidth; const y = pointsBuffer[i + 1] * image.rows * ratioHeight; points.push({ x, y }); } return new PFLDResult(points); } }