autotagger v2 lol

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
+    };
+  
+  };
 
 };

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(', ')}` : '')
         ]);
 
       });