Khrisna commited on
Commit
3dd73af
·
1 Parent(s): 7e84ef0

Update index.js

Browse files
Files changed (1) hide show
  1. index.js +15 -3
index.js CHANGED
@@ -337,9 +337,21 @@ async function waifu2x(image, formats) {
337
  }
338
  async function check_nsfw(buffer) {
339
  const model = await nsfwjs.load() // To load a local model, nsfw.load('file://./path/to/model/')
340
- // Image must be in tf.tensor3d format
341
- // you can convert image to tf.tensor3d with tf.node.decodeImage(Uint8Array,channels)
342
- const image = await tfjs.node.decodeImage(buffer,3)
 
 
 
 
 
 
 
 
 
 
 
 
343
  const predictions = await model.classify(image)
344
  image.dispose();
345
  const results = predictions.map(v => {
 
337
  }
338
  async function check_nsfw(buffer) {
339
  const model = await nsfwjs.load() // To load a local model, nsfw.load('file://./path/to/model/')
340
+ const convert = async (img) => {
341
+ // Decoded image in UInt8 Byte array
342
+ const image = await jpeg.decode(img, { useTArray: true })
343
+
344
+ const numChannels = 3
345
+ const numPixels = image.width * image.height
346
+ const values = new Int32Array(numPixels * numChannels)
347
+
348
+ for (let i = 0; i < numPixels; i++)
349
+ for (let c = 0; c < numChannels; ++c)
350
+ values[i * numChannels + c] = image.data[i * 4 + c]
351
+
352
+ return tf.tensor3d(values, [image.height, image.width, numChannels], 'int32')
353
+ },3)
354
+ const image = await convert(buffer)
355
  const predictions = await model.classify(image)
356
  image.dispose();
357
  const results = predictions.map(v => {