70 lines
3.1 KiB
TypeScript
70 lines
3.1 KiB
TypeScript
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);
|
|
}
|
|
|
|
}
|