Introducing trainer. Trainer creates one labeled facedescriptor per folder

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";
+import { Trainer } from "../src/trainer";
-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;
+  const trainer = new Trainer(
-  await faceDetectionNet.loadFromDisk(path.join(__dirname, "../weights"));
+    process.env.REF_IMAGE_DIR as string,
-  await faceapi.nets.faceLandmark68Net.loadFromDisk(
+    process.env.TRAINED_MODEL_DIR as string
-    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`);
-    }
   );
+  await trainer.train(true);
 };
 
 main();

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 {
+import { FaceMatcher } from "@vladmandic/face-api";
-  LabeledFaceDescriptors,
-  TNetInput,
-  FaceMatcher,
-} from "@vladmandic/face-api";
-import * as mime from "mime-types";
 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() {
-    //Train facial recognition model
-    let faceMatcher: FaceMatcher;
-    if (this.config.trainOnStartup) {
-      faceMatcher = await this.trainModels();
-    } 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);
 
+    //Train facial recognition model
+    let faceMatcher: FaceMatcher;
+    if (this.config.trainOnStartup) {
+      const trainer = new Trainer(
+        this.config.refImageDirectory,
+        this.config.trainedModelDirectory
+      );
+      faceMatcher = await trainer.train(true);
+    } else {
       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;
-  }
 }

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

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<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)
+          );
+        }
+      })
+    );
+  };
+}