arthur/xiao
1
0
Fork 0
mirror of https://github.com/arthur-pbty/xiao.git synced 2026-06-12 08:14:47 +02:00
Code Issues Packages Projects Releases Wiki Activity

More accurate check

Browse Source
This commit is contained in:
Dragon Fire
2021-03-22 21:43:11 -04:00
parent 832595c7d0
commit 0d8e8fd4d9
2 changed files with 17 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
commands/analyze/nsfw.js
+7 -3
View File
@@ -1,5 +1,6 @@
const Command = require('../../structures/Command'); const Command = require('../../structures/Command');
const request = require('node-superfetch'); const request = require('node-superfetch');
const { stripIndents } = require('common-tags');
const { isImageNSFW } = require('../../util/Util'); const { isImageNSFW } = require('../../util/Util');
module.exports = class NsfwCommand extends Command { module.exports = class NsfwCommand extends Command {
@@ -27,9 +28,12 @@ module.exports = class NsfwCommand extends Command {
async run(msg, { image }) { async run(msg, { image }) {
try { try {
const { body } = await request.get(image); const { body } = await request.get(image);
const prediction = await isImageNSFW(this.client.nsfwModel, body, false); const predictions = await isImageNSFW(this.client.nsfwModel, body, false);
const prob = Math.round(prediction.probability * 100); const formatted = predictions.map(result => `${Math.round(result.probability * 100)}% ${result.className}`);
return msg.reply(`I'm **${prob}%** sure this image is: **${prediction.className}**.`); return msg.reply(stripIndents`
**This image gives the following results:**
${formatted.join('\n')}
`);
} catch (err) { } catch (err) {
return msg.reply(`Oh no, an error occurred: \`${err.message}\`. Try again later!`); return msg.reply(`Oh no, an error occurred: \`${err.message}\`. Try again later!`);
} }
util/Util.js
+10 -2
View File
@@ -219,9 +219,17 @@ module.exports = class Util {
static async isImageNSFW(model, image, bool = true) { static async isImageNSFW(model, image, bool = true) {
const img = await tf.node.decodeImage(image, 3); const img = await tf.node.decodeImage(image, 3);
const predictions = await model.classify(img, 1); const predictions = await model.classify(img, 2);
img.dispose(); img.dispose();
return bool ? predictions[0].className !== 'Neutral' && predictions[0].className !== 'Drawing' : predictions[0]; const results = [];
results.push(predictions[0]);
for (const result of predictions) {
if (result.className === predictions[0].className) continue;
if (result.probability >= predictions[0].probability - 5) results.push(result);
}
return bool
? results.some(result => result.className !== 'Drawing' && result.className !== 'Neutral')
: results;
} }
static async reactIfAble(msg, user, emoji, fallbackEmoji) { static async reactIfAble(msg, user, emoji, fallbackEmoji) {