-
Notifications
You must be signed in to change notification settings - Fork 22
/
index.js
67 lines (56 loc) · 1.64 KB
/
index.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
const express = require('express')
const multer = require('multer')
const jpeg = require('jpeg-js')
const tf = require('@tensorflow/tfjs-node')
const nsfw = require('nsfwjs')
const app = express()
const sharp = require('sharp');
const upload = multer({
fileFilter: function (req, file, cb) {
if (
file.mimetype !== 'image/png' &&
file.mimetype !== 'image/jpg' &&
file.mimetype !== 'image/jpeg'
) {
cb(null, false);
} else {
cb(null, true);
}
}
});
let _model
const convert = async (img) => {
// Decoded image in UInt8 Byte array
const image = await jpeg.decode(img, true)
const numChannels = 3
const numPixels = image.width * image.height
const values = new Int32Array(numPixels * numChannels)
for (let i = 0; i < numPixels; i++)
for (let c = 0; c < numChannels; ++c)
values[i * numChannels + c] = image.data[i * 4 + c]
return tf.tensor3d(values, [image.height, image.width, numChannels], 'int32')
}
app.post('/api', upload.single('image'), async (req, res) => {
if (!req.file) res.status(400).send('Missing image multipart/form-data')
else {
try {
const data = await sharp(req.file.buffer)
.jpeg()
.toBuffer()
const image = await convert(data)
const predictions = await _model.classify(image)
image.dispose()
res.json(predictions)
} catch (error) {
res.json(400).send('错误请求!');
}
}
})
app.get('/index', async (req, res)=> {
res.send('23124124');
})
const load_model = async () => {
_model = await nsfw.load()
}
// Keep the model in memory, make sure it's loaded only once
load_model().then(() => app.listen(3027))