-
Notifications
You must be signed in to change notification settings - Fork 4
/
index.js
80 lines (65 loc) · 2.26 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
68
69
70
71
72
73
74
75
76
77
78
79
80
const express = require('express');
const axios = require('axios');
const nsfwjs = require('nsfwjs');
const tf = require('@tensorflow/tfjs-node');
const sharp = require('sharp');
const app = express();
const PORT = 9740;
const HOST = '127.0.0.1';
let model;
nsfwjs.load("https://raw.githubusercontent.com/infinitered/nsfwjs/master/models/inception_v3/model.json", { size: 299 })
.then((loadedModel) => {
model = loadedModel;
console.log('Model loaded');
})
.catch((error) => {
console.error('Error loading model:', error);
});
// Redirect root path to GitHub
app.get('/', (req, res) => {
res.redirect('https://github.com/0-RTT/nsfw');
});
// NSFW classification endpoint
app.get('/nsfw', async (req, res) => {
try {
const { url } = req.query;
if (!url) {
return res.status(400).json({ message: 'Invalid URL.' });
}
const response = await axios.get(url, { responseType: 'arraybuffer' });
const imageBuffer = Buffer.from(response.data);
let imageTensor;
if (url.endsWith('.gif')) {
const jpgBuffer = await sharp(imageBuffer)
.resize({ width: 299, height: 299 })
.toFormat('jpeg')
.toBuffer();
imageTensor = tf.node.decodeImage(jpgBuffer, 3);
} else {
imageTensor = tf.node.decodeImage(imageBuffer, 3);
}
const predictions = await model.classify(imageTensor);
imageTensor.dispose();
const formattedPredictions = predictions.reduce((acc, { className, probability }) => {
acc[className] = probability;
return acc;
}, {});
res.json(formattedPredictions);
} catch (error) {
console.error('Error processing image:', error);
if (error.response) {
return res.status(error.response.status).json({ message: 'Error fetching image from URL.', details: error.message });
} else if (error.code === 'ERR_INVALID_URL') {
return res.status(400).json({ message: 'Invalid image URL.', details: error.message });
} else {
return res.status(500).json({ message: 'Internal server error.', details: error.message });
}
}
});
// Handle 404 for any other paths
app.use((req, res) => {
res.sendStatus(404); // 直接返回 404 响应
});
app.listen(PORT, HOST, () => {
console.log(`Server is running on http://${HOST}:${PORT}`);
});