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

Training into a singular dataset file

Browse Source
This commit is contained in:
watsonb8 2020-11-27 01:15:46 -05:00
parent 53c1b162ae
commit 6268efe517
2 changed files with 39 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
scripts/streamAndDetect.ts
View File

@ -1,5 +1,5 @@
import { Rtsp } from "rtsp-stream/lib"; import { Rtsp } from "rtsp-stream/lib";
import { FaceMatcher, nets } from "@vladmandic/face-api"; import { nets } from "@vladmandic/face-api";
import * as faceapi from "@vladmandic/face-api"; import * as faceapi from "@vladmandic/face-api";
import canvas from "canvas"; import canvas from "canvas";
import fs from "fs"; import fs from "fs";
@ -31,28 +31,23 @@ const main = async () => {
path.join(__dirname, "../weights") path.join(__dirname, "../weights")
); );
const files = fs.readdirSync(modelDir); const raw = fs.readFileSync(path.join(modelDir, "data.json"), "utf-8");
const matchers: Array<FaceMatcher> = []; const content = JSON.parse(raw);
for (const file of files) { const matcher = faceapi.FaceMatcher.fromJSON(content);
const raw = fs.readFileSync(path.join(modelDir, file), "utf-8");
const content = JSON.parse(raw);
matchers.push(FaceMatcher.fromJSON(content));
}
rtsp.on("data", async (data: Buffer) => { rtsp.on("data", async (data: Buffer) => {
const img = new Image(); const input = ((await canvas.loadImage(data)) as unknown) as ImageData;
img.src = data.toString("base64"); const out = faceapi.createCanvasFromMedia(input);
const input = await canvas.loadImage(data, "base64");
// fs.writeFileSync(path.join(__dirname, "image.jpg"), data, "base64");
const resultsQuery = await faceapi const resultsQuery = await faceapi
.detectAllFaces(input, getFaceDetectorOptions(faceDetectionNet)) .detectAllFaces(out, getFaceDetectorOptions(faceDetectionNet))
.withFaceLandmarks() .withFaceLandmarks()
.withFaceDescriptors(); .withFaceDescriptors();
for (const res of resultsQuery) { for (const res of resultsQuery) {
for (const matcher of matchers) { const bestMatch = matcher.findBestMatch(res.descriptor);
const bestMatch = matcher.findBestMatch(res.descriptor); console.log(bestMatch.label);
console.log(bestMatch.label);
}
} }
}); });

50
scripts/train.ts
View File

@ -2,7 +2,7 @@ import * as faceapi from "@vladmandic/face-api";
import canvas from "canvas"; import canvas from "canvas";
import fs, { lstatSync } from "fs"; import fs, { lstatSync } from "fs";
import * as path from "path"; import * as path from "path";
import { TNetInput } from "@vladmandic/face-api"; import { LabeledFaceDescriptors, TNetInput } from "@vladmandic/face-api";
import * as mime from "mime-types"; import * as mime from "mime-types";
import dotenv from "dotenv-extended"; import dotenv from "dotenv-extended";
import { getFaceDetectorOptions } from "../src/common"; import { getFaceDetectorOptions } from "../src/common";
@ -33,6 +33,7 @@ const main = async () => {
const dirs = fs.readdirSync(inputDir); const dirs = fs.readdirSync(inputDir);
const refs: Array<LabeledFaceDescriptors> = [];
for (const dir of dirs) { for (const dir of dirs) {
if (!lstatSync(path.join(inputDir, dir)).isDirectory()) { if (!lstatSync(path.join(inputDir, dir)).isDirectory()) {
continue; continue;
@ -54,11 +55,15 @@ const main = async () => {
)) as unknown; )) as unknown;
const descriptor = await faceapi const descriptor = await faceapi
.detectAllFaces(referenceImage as TNetInput, options) .detectSingleFace(referenceImage as TNetInput, options)
.withFaceLandmarks() .withFaceLandmarks()
.withFaceDescriptors(); .withFaceDescriptor();
if (!descriptor || !descriptor.descriptor) {
throw new Error("No face found");
}
return descriptor.length > 0 ? descriptor : undefined; const faceDescriptors = [descriptor.descriptor];
return new faceapi.LabeledFaceDescriptors(dir, faceDescriptors);
} catch (err) { } catch (err) {
console.log( console.log(
"An error occurred loading image at path: " + "An error occurred loading image at path: " +
@ -69,26 +74,27 @@ const main = async () => {
}) })
); );
const items = []; if (referenceResults) {
for (const item of referenceResults) { refs.push(
if (item) { ...(referenceResults.filter((e) => e) as LabeledFaceDescriptors[])
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`);
}
);
} }
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`);
}
);
}; };
main(); main();