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();
 })();

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