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
Show Stats 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();

112
src/homeLocationPlatform.ts
View File

@ -10,17 +10,12 @@ import {
import { IConfig, isConfig } from "./config";
import * as faceapi from "@vladmandic/face-api";
import canvas from "canvas";
import fs, { lstatSync } from "fs";
import fs from "fs";
import * as path from "path";
import { nets } from "@vladmandic/face-api";
import {
LabeledFaceDescriptors,
TNetInput,
FaceMatcher,
} from "@vladmandic/face-api";
import * as mime from "mime-types";
import { FaceMatcher } from "@vladmandic/face-api";
import { Monitor } from "./monitor/monitor";
import { getFaceDetectorOptions } from "./common";
import { Trainer } from "./trainer";
require("@tensorflow/tfjs-node");
const { Canvas, Image, ImageData } = canvas;
@ -82,17 +77,20 @@ export class HomeLocationPlatform implements DynamicPlatformPlugin {
* must not be registered again to prevent "duplicate UUID" errors.
*/
public async discoverDevices() {
const faceDetectionNet = nets.ssdMobilenetv1;
await faceDetectionNet.loadFromDisk(this.config.weightDirectory);
await nets.faceLandmark68Net.loadFromDisk(this.config.weightDirectory);
await nets.faceRecognitionNet.loadFromDisk(this.config.weightDirectory);
//Train facial recognition model
let faceMatcher: FaceMatcher;
if (this.config.trainOnStartup) {
faceMatcher = await this.trainModels();
const trainer = new Trainer(
this.config.refImageDirectory,
this.config.trainedModelDirectory
);
faceMatcher = await trainer.train(true);
} else {
const faceDetectionNet = nets.ssdMobilenetv1;
await faceDetectionNet.loadFromDisk(this.config.weightDirectory);
await nets.faceLandmark68Net.loadFromDisk(this.config.weightDirectory);
await nets.faceRecognitionNet.loadFromDisk(this.config.weightDirectory);
const raw = fs.readFileSync(
path.join(this.config.trainedModelDirectory, "data.json"),
"utf-8"
@ -142,88 +140,4 @@ export class HomeLocationPlatform implements DynamicPlatformPlugin {
}
}
}
private async trainModels(): Promise<FaceMatcher> {
const faceDetectionNet = faceapi.nets.ssdMobilenetv1;
await faceDetectionNet.loadFromDisk(this.config.weightDirectory);
await faceapi.nets.faceLandmark68Net.loadFromDisk(
this.config.weightDirectory
);
await faceapi.nets.faceRecognitionNet.loadFromDisk(
this.config.weightDirectory
);
const options = getFaceDetectorOptions(faceDetectionNet);
const dirs = fs.readdirSync(this.config.refImageDirectory);
const refs: Array<LabeledFaceDescriptors> = [];
for (const dir of dirs) {
if (
!lstatSync(path.join(this.config.refImageDirectory, dir)).isDirectory()
) {
continue;
}
const files = fs.readdirSync(
path.join(this.config.refImageDirectory, dir)
);
let referenceResults = await Promise.all(
files.map(async (file: string) => {
const mimeType = mime.contentType(
path.extname(path.join(this.config.refImageDirectory, dir, file))
);
if (!mimeType || !mimeType.startsWith("image")) {
return;
}
this.log.info(path.join(this.config.refImageDirectory, dir, file));
try {
const referenceImage = (await canvas.loadImage(
path.join(this.config.refImageDirectory, 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) {
this.log.info(
"An error occurred loading image at path: " +
path.join(this.config.refImageDirectory, dir, file)
);
}
return undefined;
})
);
if (referenceResults) {
refs.push(
...(referenceResults.filter((e) => e) as LabeledFaceDescriptors[])
);
}
}
const faceMatcher = new faceapi.FaceMatcher(refs);
fs.writeFile(
path.join(this.config.trainedModelDirectory, "data.json"),
JSON.stringify(faceMatcher.toJSON()),
"utf8",
(err) => {
if (err) {
this.log.info(`An error occurred while writing data model to file`);
}
this.log.info(`Successfully wrote data model to file`);
}
);
return faceMatcher;
}
}

4
src/monitor/monitor.ts
View File

@ -208,7 +208,9 @@ export class Monitor {
private onWatchdogTimeout = async (stream: IStream, roomName: string) => {
this._logger.info(
`[${stream.connectionString}] Watchdog timeout: restarting stream`
`[${this.getRedactedConnectionString(
stream.connectionString
)}] Watchdog timeout: restarting stream`
);
//Close and remove old stream

103
src/trainer.ts Normal file
View File

@ -0,0 +1,103 @@
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 { getFaceDetectorOptions } from "./common";
require("@tensorflow/tfjs-node");
const { Canvas, Image, ImageData } = canvas;
//@ts-ignore
faceapi.env.monkeyPatch({ Canvas, Image, ImageData });
export class Trainer {
constructor(private _refImageDir: string, private _trainedModelDir: string) {}
public async train(writeToDisk: boolean): Promise<faceapi.FaceMatcher> {
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(this._refImageDir);
const refs = [];
for (const dir of dirs) {
const descriptor = new LabeledFaceDescriptors(dir, []);
await this.getLabeledFaceDescriptorFromDir(
path.join(this._refImageDir, dir),
descriptor,
options
);
if (descriptor) {
refs.push(descriptor);
}
}
const faceMatcher = new faceapi.FaceMatcher(refs);
if (writeToDisk) {
fs.writeFile(
path.join(this._trainedModelDir, "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`);
}
);
}
return faceMatcher;
}
private getLabeledFaceDescriptorFromDir = async (
dir: string,
labeldFaceDescriptors: LabeledFaceDescriptors,
options: faceapi.TinyFaceDetectorOptions | faceapi.SsdMobilenetv1Options
): Promise<void> => {
if (!lstatSync(dir).isDirectory()) {
return;
}
const files = fs.readdirSync(dir);
await Promise.all(
files.map(async (file: string) => {
const mimeType = mime.contentType(path.extname(path.join(dir, file)));
if (!mimeType || !mimeType.startsWith("image")) {
return;
}
console.log(path.join(dir, file));
try {
const referenceImage = (await canvas.loadImage(
path.join(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");
}
labeldFaceDescriptors.descriptors.push(descriptor.descriptor);
} catch (err) {
throw new Error(
"An error occurred loading image at path: " + path.join(dir, file)
);
}
})
);
};
}