101 lines
2.9 KiB
TypeScript
101 lines
2.9 KiB
TypeScript
import * as faceapi from "@vladmandic/face-api";
|
|
import canvas from "canvas";
|
|
import fs, { lstatSync } from "fs";
|
|
import * as path from "path";
|
|
import { LabeledFaceDescriptors, TNetInput } from "@vladmandic/face-api";
|
|
import * as mime from "mime-types";
|
|
import dotenv from "dotenv-extended";
|
|
import { getFaceDetectorOptions } from "../src/common";
|
|
require("@tensorflow/tfjs-node");
|
|
|
|
const { Canvas, Image, ImageData } = canvas;
|
|
//@ts-ignore
|
|
faceapi.env.monkeyPatch({ Canvas, Image, ImageData });
|
|
|
|
const main = async () => {
|
|
dotenv.load({
|
|
silent: false,
|
|
errorOnMissing: true,
|
|
});
|
|
const inputDir = process.env.REF_IMAGE_DIR as string;
|
|
const outDir = process.env.TRAINED_MODEL_DIR as string;
|
|
|
|
const faceDetectionNet = faceapi.nets.ssdMobilenetv1;
|
|
await faceDetectionNet.loadFromDisk(path.join(__dirname, "../weights"));
|
|
await faceapi.nets.faceLandmark68Net.loadFromDisk(
|
|
path.join(__dirname, "../weights")
|
|
);
|
|
await faceapi.nets.faceRecognitionNet.loadFromDisk(
|
|
path.join(__dirname, "../weights")
|
|
);
|
|
|
|
const options = getFaceDetectorOptions(faceDetectionNet);
|
|
|
|
const dirs = fs.readdirSync(inputDir);
|
|
|
|
const refs: Array<LabeledFaceDescriptors> = [];
|
|
for (const dir of dirs) {
|
|
if (!lstatSync(path.join(inputDir, dir)).isDirectory()) {
|
|
continue;
|
|
}
|
|
const files = fs.readdirSync(path.join(inputDir, dir));
|
|
let referenceResults = await Promise.all(
|
|
files.map(async (file: string) => {
|
|
const mimeType = mime.contentType(
|
|
path.extname(path.join(inputDir, dir, file))
|
|
);
|
|
if (!mimeType || !mimeType.startsWith("image")) {
|
|
return;
|
|
}
|
|
console.log(path.join(inputDir, dir, file));
|
|
|
|
try {
|
|
const referenceImage = (await canvas.loadImage(
|
|
path.join(inputDir, dir, file)
|
|
)) as unknown;
|
|
|
|
const descriptor = await faceapi
|
|
.detectSingleFace(referenceImage as TNetInput, options)
|
|
.withFaceLandmarks()
|
|
.withFaceDescriptor();
|
|
if (!descriptor || !descriptor.descriptor) {
|
|
throw new Error("No face found");
|
|
}
|
|
|
|
const faceDescriptors = [descriptor.descriptor];
|
|
return new faceapi.LabeledFaceDescriptors(dir, faceDescriptors);
|
|
} catch (err) {
|
|
console.log(
|
|
"An error occurred loading image at path: " +
|
|
path.join(inputDir, dir, file)
|
|
);
|
|
}
|
|
return undefined;
|
|
})
|
|
);
|
|
|
|
if (referenceResults) {
|
|
refs.push(
|
|
...(referenceResults.filter((e) => e) as LabeledFaceDescriptors[])
|
|
);
|
|
}
|
|
}
|
|
|
|
const faceMatcher = new faceapi.FaceMatcher(refs);
|
|
|
|
fs.writeFile(
|
|
path.join(outDir, "data.json"),
|
|
JSON.stringify(faceMatcher.toJSON()),
|
|
"utf8",
|
|
(err) => {
|
|
if (err) {
|
|
console.log(`An error occurred while writing data model to file`);
|
|
}
|
|
|
|
console.log(`Successfully wrote data model to file`);
|
|
}
|
|
);
|
|
};
|
|
|
|
main();
|