diff --git a/debug/autotagger.mjs b/debug/autotagger.mjs index 0b6d293..d2bc205 100644 --- a/debug/autotagger.mjs +++ b/debug/autotagger.mjs @@ -1,33 +1,60 @@ -import sql from "../src/inc/sql.mjs"; -import cfg from "../src/inc/config.mjs"; -import fetch from "flumm-fetch-cookies"; +import db from "../src/inc/sql.mjs"; +import lib from "../src/inc/lib.mjs"; (async () => { const _args = process.argv.slice(2); const _from = +_args[0]; - const _to = _from + 100; + const _to = _from + 500; - const f0cks = await sql("items") - .whereRaw("id not in (select item_id from tags_assign group by item_id)") - .andWhere("mime", "like", "image/%") - .andWhereBetween("id", [ _from, _to ]); + const f0cks = await db` + select * + from items + where + id not in (select item_id from tags_assign group by item_id) and + mime like 'image/%' and + id between ${_from} and ${_to} + `; - const body = { headers: { Authorization: `Basic ${cfg.tagger.btoa}` } }; - + console.time('blah'); for(let f of f0cks) { - let tag; - const res = await (await fetch(`${cfg.tagger.endpoint}/categories/nsfw_beta?image_url=${cfg.main.url}/b/${f.dest}`, body)).json(); - if(res?.result) { - tag = (res.result.categories[0].name.en === "safe") ? "sfw" : "nsfw"; + const tmp = await lib.detectNSFW(f.dest); - await sql("tags_assign").insert({ - tag_id: tag === "sfw" ? 1 : 2, - item_id: f.id, - user_id: 7 // user: autotagger - }); - } - else { - console.log(res); + console.log( + 'https://f0ck.me/' + f.id, + tmp.isNSFW, + tmp.score.toFixed(2), + { + sexy: tmp.scores.sexy.toFixed(2), + porn: tmp.scores.porn.toFixed(2), + hentai: tmp.scores.hentai.toFixed(2), + neutral: tmp.scores.neutral.toFixed(2) + } + ); + + await db` + insert into "tags_assign" ${ + db({ + item_id: f.id, + tag_id: tmp.nsfw ? 2 : 1, + user_id: 7 + }) + } + `; + + if(tmp.hentai >= .7) { + await db` + insert into "tags_assign" ${ + db({ + item_id: f.id, + tag_id: 8, // hentai + user_id: 7 // autotagger + }) + } + `; } + }; + + console.timeEnd('blah'); + process.exit(); })(); diff --git a/nsfw_model.h5 b/nsfw_model.h5 new file mode 100644 index 0000000..6b2d2a1 Binary files /dev/null and b/nsfw_model.h5 differ diff --git a/src/inc/lib.mjs b/src/inc/lib.mjs index 2cf7214..ff6bfea 100644 --- a/src/inc/lib.mjs +++ b/src/inc/lib.mjs @@ -1,7 +1,9 @@ import crypto from "crypto"; import util from "util"; import db from "./sql.mjs"; +import { exec as _exec } from "child_process"; +const exec = util.promisify(_exec); const scrypt = util.promisify(crypto.scrypt); const epochs = [ @@ -150,5 +152,31 @@ export default new class { } return tags; }; + async detectNSFW(dest) { + const { stdout, stderr } = await exec( + `python -c "import sys\nfrom nsfw_detector import predict\nmodel = predict.load_model('./nsfw_model.h5')\nprint(predict.classify(model, './public/b/${dest}'))"` + ); + const res = JSON.parse(stdout.replace(/\'/g, '"').split('\n').slice(1, -1)); + const tmp = Object.values(res)[0]; + + tmp.sexy = tmp.sexy / 2; + + let nsfw = false; + if(tmp.neutral >= .7) + nsfw = false; + else if((tmp.sexy + tmp.porn + tmp.hentai) >= .7) + nsfw = true; + else if(tmp.drawings >= .4) + nsfw = false; + else + nsfw = false; + + return { + isNSFW: nsfw, + score: tmp.sexy + tmp.porn + tmp.hentai, + scores: tmp + }; + + }; }; diff --git a/src/inc/trigger/parser.mjs b/src/inc/trigger/parser.mjs index 65036d1..46ecf26 100644 --- a/src/inc/trigger/parser.mjs +++ b/src/inc/trigger/parser.mjs @@ -191,35 +191,57 @@ export default async bot => { speed = !Number.isFinite(speed) ? "yes" : `${speed.toFixed(2)} Mbit/s`; // autotagger - let tag; + let tags = []; try { - if(mime.startsWith('image') && mime != 'image/gif') { - const body = { headers: { Authorization: `Basic ${cfg.tagger.btoa}` } }; - const res = await (await fetch(`${cfg.tagger.endpoint}/categories/nsfw_beta?image_url=${cfg.main.url}/b/${filename}`, body)).json(); - if(res) - tag = (res.result.categories[0].name.en === 'safe') ? 'sfw' : 'nsfw'; - } - else if(mime.startsWith('audio')) { - tag = 'sfw'; - } + if(mime.startsWith('image')) { + const res = await lib.detectNSFW(filename); - if(tag === 'sfw' || tag === 'nsfw') { await db` insert into "tags_assign" ${ db({ - tag_id: tag === "sfw" ? 1 : 2, item_id: itemid, - user_id: 7 // user: autotagger (ID: 7) - }, 'tag_id', 'item_id', 'user_id') + tag_id: res.nsfw ? 2 : 1, + user_id: 7 + }) } `; + tags.push(res.nsfw ? 'nsfw' : 'sfw'); + + if(res.hentai >= .7) { + await db` + insert into "tags_assign" ${ + db({ + item_id: f.id, + tag_id: 8, // hentai + user_id: 7 // autotagger + }) + } + `; + tags.push('hentai'); + } + } + else if(mime.startsWith('audio')) { + await db` + insert into "tags_assign" ${ + db([{ + item_id: itemid, + tag_id: 1, + user_id: 7 + }, { + item_id: itemid, + tag_id: 7178, + user_id: 7 + }]) + } + `; + tags.push('sfw', 'audio'); } } catch(err) { console.error(err); } e.reply([ - `[f0cked] link: ${cfg.main.url}/${itemid} | size: ${lib.formatSize(size)} | speed: ${speed}` + (tag ? ` | ${tag}` : "") + `[f0cked] link: ${cfg.main.url}/${itemid} | size: ${lib.formatSize(size)} | speed: ${speed}` + (tags.length > 0 ? ` | tags: ${tags.join(', ')}` : '') ]); });