初次提交

src/deploy/facealign/pfld.ts

@@ -0,0 +1,53 @@
+import { writeFileSync } from "fs";
+import { cv } from "../../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 {
+	protected directionPointIndex(): [number, number] { return [36, 92]; }
+	protected leftEyePointIndex(): number[] { return [33, 34, 35, 36, 37, 38, 39, 40, 41, 42]; }
+	protected rightEyePointIndex(): number[] { return [87, 88, 89, 90, 91, 92, 93, 94, 95, 96]; }
+	protected leftEyebrowPointIndex(): number[] { return [43, 44, 45, 46, 47, 48, 49, 50, 51]; }
+	protected rightEyebrowPointIndex(): number[] { return [97, 98, 99, 100, 101, 102, 103, 104, 105]; }
+	protected mouthPointIndex(): number[] { return [52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71]; }
+	protected nosePointIndex(): number[] { return [72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86]; }
+	protected contourPointIndex(): number[] { return [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]; }
+}
+
+export class PFLD extends 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 = image.crop(option.crop.sx, option.crop.sy, option.crop.sw, option.crop.sh);
+		const ratioWidth = image.width / input.shape[3];
+		const ratioHeight = image.height / input.shape[2];
+		image = image.resize(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];
+
+		const points: FacePoint[] = [];
+		for (let i = 0; i < pointsBuffer.length; i += 2) {
+			const x = pointsBuffer[i] * image.width * ratioWidth;
+			const y = pointsBuffer[i + 1] * image.height * ratioHeight;
+			points.push({ x, y });
+		}
+
+		return new PFLDResult(points);
+	}
+
+}