109 lines
3.1 KiB
TypeScript
109 lines
3.1 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";
|
|
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);
|
|
|
|
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
|
|
.detectAllFaces(referenceImage as TNetInput, options)
|
|
.withFaceLandmarks()
|
|
.withFaceDescriptors();
|
|
|
|
return descriptor.length > 0 ? descriptor : undefined;
|
|
} catch (err) {
|
|
console.log(
|
|
"An error occurred loading image at path: " +
|
|
path.join(inputDir, dir, file)
|
|
);
|
|
}
|
|
return undefined;
|
|
})
|
|
);
|
|
|
|
const items = [];
|
|
for (const item of referenceResults) {
|
|
if (item) {
|
|
items.push(...item);
|
|
}
|
|
}
|
|
const faceMatcher = new faceapi.FaceMatcher(items);
|
|
fs.writeFile(
|
|
path.join(outDir, dir + ".json"),
|
|
JSON.stringify(faceMatcher.toJSON()),
|
|
"utf8",
|
|
(err) => {
|
|
if (err) {
|
|
console.log(`An error occurred while writing ${dir} model to file`);
|
|
}
|
|
|
|
console.log(`Successfully wrote ${dir} model to file`);
|
|
}
|
|
);
|
|
}
|
|
};
|
|
|
|
// SsdMobilenetv1Options
|
|
const minConfidence = 0.5;
|
|
|
|
// TinyFaceDetectorOptions
|
|
const inputSize = 408;
|
|
const scoreThreshold = 0.5;
|
|
|
|
function getFaceDetectorOptions(net: faceapi.NeuralNetwork<any>) {
|
|
return net === faceapi.nets.ssdMobilenetv1
|
|
? new faceapi.SsdMobilenetv1Options({ minConfidence })
|
|
: new faceapi.TinyFaceDetectorOptions({ inputSize, scoreThreshold });
|
|
}
|
|
|
|
const baseDir = path.resolve(__dirname, "../out");
|
|
|
|
main();
|