watsonb8/homebridge-face-location
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Introducing trainer. Trainer creates one labeled facedescriptor per folder

Browse Source
This commit is contained in:
watsonb8
2020-12-12 13:44:51 -05:00
parent 0cb3862843
commit cf2c20b6e0
4 changed files with 124 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
scripts/train.ts
View File

@ -1,100 +1,16 @@
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 });
import { Trainer } from "../src/trainer";
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`);
}
const trainer = new Trainer(
process.env.REF_IMAGE_DIR as string,
process.env.TRAINED_MODEL_DIR as string
);
await trainer.train(true);
};
main();