-
Notifications
You must be signed in to change notification settings - Fork 0
/
face-vectorization.js
106 lines (84 loc) · 3.09 KB
/
face-vectorization.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
module.exports = function (RED) {
const onnx = require("onnxruntime-node");
const sharp = require("sharp");
const path = require("path");
const fs = require("fs");
function FaceVectorizationNode(config) {
RED.nodes.createNode(this, config);
const node = this;
const currentDir = __dirname;
const modelPath = path.join(currentDir, "model", "facenet-model.onnx");
node.on("input", async function (msg) {
const inputData = msg.payload;
try {
const vectors = [];
if (config.inputType == 0) {
vectors.push(await imageVectorization(inputData));
} else if (config.inputType == 1) {
for (let i = 0; i < inputData.length; ++i) {
vectors.push(await imageVectorization(inputData[i]));
}
}
if (config.returnType == 0) {
msg.payload = vectors;
node.send(msg);
} else if (config.returnType == 1) {
const filePath = config.path;
const textData = JSON.stringify(vectors);
const dir = path.dirname(filePath);
if (!fs.existsSync(dir)) {
fs.mkdirSync(dir, { recursive: true });
}
if (config.method == 0) {
fs.writeFileSync(filePath, textData, "utf-8");
} else if (config.method == 1) {
let existData = fs.readFileSync(filePath, "utf-8");
const allData = JSON.parse(existData);
allData.push(...vectors);
fs.writeFileSync(filePath, JSON.stringify(allData), "utf-8");
}
}
} catch (error) {
node.error("An error occurred: " + error.message.split("\n")[0]);
}
});
async function imageVectorization(inputData) {
try {
const model = await onnx.InferenceSession.create(modelPath);
const inputName = model.inputNames[0];
const inputDims = [1, 3, 224, 224];
const input = await imageTransfer(inputData);
const inputTensor = new onnx.Tensor(Float32Array.from(input), inputDims);
const feeds = {};
feeds[inputName] = inputTensor;
const outputData = await model.run(feeds);
return Array.from(outputData.output.data);
} catch (error) {
throw new Error("Image vectorization error: " + error.message.split("\n")[0]);
}
}
async function imageTransfer(inputData) {
try {
const img = sharp(inputData);
const pixels = await img
.removeAlpha()
.resize({ width: 224, height: 224, fit: "fill" })
.raw()
.toBuffer();
const red = [],
green = [],
blue = [];
for (let index = 0; index < pixels.length; index += 3) {
red.push(pixels[index] / 255.0);
green.push(pixels[index + 1] / 255.0);
blue.push(pixels[index + 2] / 255.0);
}
const input = [...red, ...green, ...blue];
return input;
} catch (error) {
throw new Error("Image transfer error: " + error.message.split("\n")[0]);
}
}
}
RED.nodes.registerType("good-face-vectorization", FaceVectorizationNode);
};