Spaces:
Running
Running
File size: 5,033 Bytes
ebe720b |
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 |
(async function () {
require('dotenv').config()
const express = require('express')
const tf = require("@tensorflow/tfjs-node")
const sharp = require("sharp");
const jpeg = require("jpeg-js")
const ffmpeg = require("fluent-ffmpeg")
const { fileTypeFromBuffer } = (await import('file-type'));
const stream = require("stream")
const ffmpegPath = require('@ffmpeg-installer/ffmpeg').path;
const ffprobePath = require('@ffprobe-installer/ffprobe').path;
const nsfwjs = require("nsfwjs");
const fs = require("fs")
ffmpeg.setFfprobePath(ffprobePath);
ffmpeg.setFfmpegPath(ffmpegPath);
// require("./model").loadModel()
const app = express()
const model = await nsfwjs.load("InceptionV3");
app.use(express.json())
app.all('/', async (req, res) => {
try {
const { img, auth } = req.query
if (img) {
if (process.env.AUTH) {
if (!auth || process.env.AUTH != auth) return res.send("Invalid auth code")
}
const imageBuffer = await fetch(img).then(async c => await c.arrayBuffer())
// console.log((await fileTypeFromBuffer(imageBuffer)).mime)
if ((await fileTypeFromBuffer(imageBuffer)).mime.includes("image")) {
const convertedBuffer = await sharp(Buffer.from(imageBuffer)).jpeg().toBuffer(); // convert webp to jpeg
const image = await convert(convertedBuffer)
const predictions = await model.classify(image);
image.dispose(); // Tensor memory must be managed explicitly (it is not sufficient to let a tf.Tensor go out of scope for its memory to be released).
return res.send(predictions);
} else {
let inputStream1 = new stream.PassThrough();
inputStream1.end(Buffer.from(imageBuffer));
ffmpeg.ffprobe(inputStream1, function (err, metadata) {
if (err) {
console.error(err);
return;
}
// Get a random second
const randomSecond = Math.floor(Math.random() * metadata.format.duration);
// Create a new input stream for the ffmpeg command
let inputStream2 = new stream.PassThrough();
inputStream2.end(Buffer.from(imageBuffer));
// Create a PassThrough stream to collect the output
const output = new stream.PassThrough();
// Set up the ffmpeg command
ffmpeg({ source: inputStream2 })
.seekInput(randomSecond)
.outputOptions('-vframes', '1')
.outputOptions('-f', 'image2pipe')
.outputOptions('-vcodec', 'png')
.output(output)
.on('error', console.error)
.run();
// Collect the output into a buffer
const chunks = [];
output.on('data', chunk => chunks.push(chunk));
output.on('end', async () => {
const buffer = Buffer.concat(chunks);
fs.writeFileSync("aa.png", buffer)
const convertedBuffer = await sharp(buffer).jpeg().toBuffer(); // convert webp to jpeg
const cimage = await convert(convertedBuffer)
const apredictions = await model.classify(cimage);
cimage.dispose(); // Tensor memory must be managed explicitly (it is not sufficient to let a tf.Tensor go out of scope for its memory to be released).
return res.send(apredictions);
});
});
}
}else{
return res.send('Hello World!')
}
} catch (err) {
console.log(err)
return res.status(500).json({ error: err.toString() })
}
})
const port = process.env.PORT || process.env.SERVER_PORT || 7860
app.listen(port, () => {
console.log(`Example app listening on port ${port}`)
})
const convert = async (img) => {
// Decoded image in UInt8 Byte array
const image = await jpeg.decode(img, { useTArray: 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");
};
})() |