From cf2c20b6e091cec9c70b700bc29f8f4e11734813 Mon Sep 17 00:00:00 2001 From: watsonb8 Date: Sat, 12 Dec 2020 13:44:51 -0500 Subject: [PATCH] Introducing trainer. Trainer creates one labeled facedescriptor per folder --- scripts/train.ts | 94 ++---------------------------- src/homeLocationPlatform.ts | 112 +++++------------------------------- src/monitor/monitor.ts | 4 +- src/trainer.ts | 103 +++++++++++++++++++++++++++++++++ 4 files changed, 124 insertions(+), 189 deletions(-) create mode 100644 src/trainer.ts diff --git a/scripts/train.ts b/scripts/train.ts index 8a38547..468c0e3 100644 --- a/scripts/train.ts +++ b/scripts/train.ts @@ -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 = []; - 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(); diff --git a/src/homeLocationPlatform.ts b/src/homeLocationPlatform.ts index 070a704..4839976 100644 --- a/src/homeLocationPlatform.ts +++ b/src/homeLocationPlatform.ts @@ -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 { - 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 = []; - 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; - } } diff --git a/src/monitor/monitor.ts b/src/monitor/monitor.ts index c9e59ed..078e5d0 100644 --- a/src/monitor/monitor.ts +++ b/src/monitor/monitor.ts @@ -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 diff --git a/src/trainer.ts b/src/trainer.ts new file mode 100644 index 0000000..bab8a27 --- /dev/null +++ b/src/trainer.ts @@ -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 { + 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 => { + 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) + ); + } + }) + ); + }; +}