commit 7416964f4312cadf52a81c61f28c88a924c992ce
author Skyler Grey <skyler3665@gmail.com> Thu Mar 09 11:59:09 2023 +0000
committer Skyler Grey <skyler3665@gmail.com> Thu Mar 09 12:00:31 2023 +0000
tree 97f4eceaf1b751f712493a889f0b48263f49557e
parent 96fcc75c5f4dd99c23b5020b15575a4ff1fb6f8a

I believe NSFW PFP scanning now works. Please could reviewers test this using either server profiles or isolated alts

src/events/messageCreate.ts

diff --git a/src/events/messageCreate.ts b/src/events/messageCreate.ts
index 38c7674..cc69bf0 100644
--- a/src/events/messageCreate.ts
+++ b/src/events/messageCreate.ts
@@ -114,7 +114,7 @@
                 if (
                     config.filters.images.NSFW &&
                     !(message.channel instanceof ThreadChannel ? message.channel.parent?.nsfw : message.channel.nsfw) &&
-                    (await NSFWCheck(element))
+                    (await NSFWCheck(element.url))
                 ) {
                     messageException(message.guild.id, message.channel.id, message.id);
                     await message.delete();

src/reflex/scanners.ts

diff --git a/src/reflex/scanners.ts b/src/reflex/scanners.ts
index 74887d2..8b8d3c8 100644
--- a/src/reflex/scanners.ts
+++ b/src/reflex/scanners.ts
@@ -30,12 +30,8 @@
     }
 });
 
-export async function testNSFW(attachment: {
-    url: string;
-    height: number | null;
-    width: number | null;
-}): Promise<NSFWSchema> {
-    const [fileStream, hash] = await streamAttachment(attachment.url);
+export async function testNSFW(url: string): Promise<NSFWSchema> {
+    const [fileStream, hash] = await streamAttachment(url);
     const alreadyHaveCheck = await client.database.scanCache.read(hash);
     if (alreadyHaveCheck && "nsfw" in alreadyHaveCheck!) {
         return { nsfw: alreadyHaveCheck.nsfw };
@@ -49,11 +45,11 @@
                 resolve(buf);
             })
     )) as Buffer;
-    const array = new Uint8Array(converted);
 
-    const img = tf.node.decodeImage(array) as tf.Tensor3D;
+    const img = tf.node.decodeImage(converted, 3, undefined, false) as tf.Tensor3D;
 
     const predictions = (await nsfw_model.classify(img, 1))[0]!;
+    img.dispose();
     console.log(2, predictions);
 
     const nsfw = predictions.className === "Hentai" || predictions.className === "Porn";
@@ -155,13 +151,9 @@
     return detectionsTypes as string[];
 }
 
-export async function NSFWCheck(element: {
-    url: string;
-    height: number | null;
-    width: number | null;
-}): Promise<boolean> {
+export async function NSFWCheck(url: string): Promise<boolean> {
     try {
-        return (await testNSFW(element)).nsfw;
+        return (await testNSFW(url)).nsfw;
     } catch (e) {
         console.log(e);
         return false;
@@ -237,7 +229,7 @@
     console.log(4, avatarTextCheck);
     // Is the profile picture NSFW
     const avatar = member.displayAvatarURL({ extension: "png", size: 1024, forceStatic: true });
-    const avatarCheck = guildData.filters.images.NSFW && (await NSFWCheck({ url: avatar, height: 1024, width: 1024 }));
+    const avatarCheck = guildData.filters.images.NSFW && (await NSFWCheck(avatar));
     console.log(5, avatarCheck);
     // Does the username contain an invite
     const inviteCheck = guildData.filters.invite.enabled && /discord\.gg\/[a-zA-Z0-9]+/gi.test(member.user.username);

src/reflex/verify.ts

diff --git a/src/reflex/verify.ts b/src/reflex/verify.ts
index bdf4aba..540fb4d 100644
--- a/src/reflex/verify.ts
+++ b/src/reflex/verify.ts
@@ -121,11 +121,9 @@
             ]
         });
         if (
-            await NSFWCheck({
-                url: (interaction.member as GuildMember).user.displayAvatarURL({ extension: "png", forceStatic: true }),
-                height: 1024,
-                width: 1024
-            })
+            await NSFWCheck(
+                (interaction.member as GuildMember).user.displayAvatarURL({ extension: "png", forceStatic: true })
+            )
         ) {
             return await interaction.editReply({
                 embeds: [