Skip to content

Commit

Permalink
feat(automod:message_rules:nsfw): powerful AI NSFW image detection
Browse files Browse the repository at this point in the history
  • Loading branch information
@virtual-designer
virtual-designer committed Feb 18, 2024
1 parent b95ba08 commit 43eb34e
Show file tree
Hide file tree
Showing 5 changed files with 252 additions and 124 deletions.
209 changes: 106 additions & 103 deletions package.json
View file
Original file line number Diff line number Diff line change
@@ -1,104 +1,107 @@
{
"name": "sudobot",
"description": "A Discord bot for moderation purposes.",
"version": "8.12.0",
"main": "build/index.js",
"license": "GPL-3.0-or-later",
"keywords": [
"bot",
"discord-bot",
"moderation",
"automoderation",
"discord-moderation",
"discord-moderation-bot",
"discord-automoderation",
"discord-automoderation-bot"
],
"homepage": "https://github.com/onesoft-sudo/sudobot",
"icon": "https://res.cloudinary.com/rakinar2/image/upload/v1659628446/SudoBot-new_cvwphw.png",
"readme": "https://github.com/onesoft-sudo/sudobot#readme",
"author": {
"name": "Ar Rakin",
"email": "rakinar2@onesoftnet.eu.org",
"url": "https://virtual-designer.github.io"
},
"repository": {
"type": "git",
"url": "https://github.com/onesoft-sudo/sudobot"
},
"bugs": {
"url": "https://github.com/onesoft-sudo/sudobot/issues",
"email": "sudobot@onesoftnet.eu.org"
},
"dependencies": {
"@prisma/client": "^5.8.1",
"ascii-table3": "^0.8.2",
"async-mutex": "^0.4.0",
"axios": "^1.4.0",
"bcrypt": "^5.1.0",
"canvas": "^2.11.2",
"chalk": "^4.1.2",
"cors": "^2.8.5",
"date-fns": "^2.30.0",
"deepmerge": "^4.3.1",
"discord.js": "^14.13.0",
"dot-object": "^2.1.4",
"dotenv": "^16.3.1",
"express": "^4.18.2",
"express-rate-limit": "^6.9.0",
"figlet": "^1.7.0",
"googleapis": "^126.0.1",
"json5": "^2.2.3",
"jsonwebtoken": "^9.0.1",
"module-alias": "^2.2.3",
"pm2": "^5.3.1",
"reflect-metadata": "^0.1.13",
"semver": "^7.5.4",
"socket.io": "^4.7.2",
"tesseract.js": "^5.0.4",
"tslib": "^2.6.1",
"undici": "^5.23.0",
"uuid": "^9.0.0",
"zod": "^3.21.4"
},
"devDependencies": {
"@commitlint/cli": "^17.6.6",
"@commitlint/config-conventional": "^17.6.6",
"@types/bcrypt": "^5.0.0",
"@types/bun": "^1.0.5",
"@types/cors": "^2.8.13",
"@types/dot-object": "^2.1.2",
"@types/express": "^4.17.17",
"@types/figlet": "^1.5.8",
"@types/jsonwebtoken": "^9.0.2",
"@types/module-alias": "^2.0.2",
"@types/node": "^20.4.0",
"@types/semver": "^7.5.4",
"@types/uuid": "^9.0.2",
"husky": "latest",
"prisma": "^5.8.1",
"ts-node": "^10.9.1",
"typescript": "^5.3.3",
"zod-to-json-schema": "^3.21.4"
},
"optionalDependencies": {
"openai": "^4.26.0"
},
"scripts": {
"start": "node build/index.js",
"start:bun": "bun build/bun.js",
"prepare": "husky",
"dev": "bun run src/bun.ts",
"lint": "eslint src/ --ext .ts",
"lint:fix": "eslint src/ --ext .ts --fix",
"build": "tsc",
"start:prod": "pm2 start ./ecosystem.config.js",
"deploy": "node scripts/deploy-commands.js",
"gen:schema": "node scripts/generate-config-schema.js",
"clean": "rm -frv build; make clean",
"test": "node --test --require ts-node/register tests/**/*.test.ts"
},
"_moduleAliases": {
"@sudobot": "build"
}
}
"name": "sudobot",
"description": "A Discord bot for moderation purposes.",
"version": "8.12.0",
"main": "build/index.js",
"license": "GPL-3.0-or-later",
"keywords": [
"bot",
"discord-bot",
"moderation",
"automoderation",
"discord-moderation",
"discord-moderation-bot",
"discord-automoderation",
"discord-automoderation-bot"
],
"homepage": "https://github.com/onesoft-sudo/sudobot",
"icon": "https://res.cloudinary.com/rakinar2/image/upload/v1659628446/SudoBot-new_cvwphw.png",
"readme": "https://github.com/onesoft-sudo/sudobot#readme",
"author": {
"name": "Ar Rakin",
"email": "rakinar2@onesoftnet.eu.org",
"url": "https://virtual-designer.github.io"
},
"repository": {
"type": "git",
"url": "https://github.com/onesoft-sudo/sudobot"
},
"bugs": {
"url": "https://github.com/onesoft-sudo/sudobot/issues",
"email": "sudobot@onesoftnet.eu.org"
},
"dependencies": {
"@prisma/client": "^5.8.1",
"@tensorflow/tfjs-node": "^4.17.0",
"ascii-table3": "^0.8.2",
"async-mutex": "^0.4.0",
"axios": "^1.4.0",
"bcrypt": "^5.1.0",
"canvas": "^2.11.2",
"chalk": "^4.1.2",
"cors": "^2.8.5",
"date-fns": "^2.30.0",
"deepmerge": "^4.3.1",
"discord.js": "^14.14.1",
"dot-object": "^2.1.4",
"dotenv": "^16.3.1",
"express": "^4.18.2",
"express-rate-limit": "^6.9.0",
"figlet": "^1.7.0",
"googleapis": "^126.0.1",
"json5": "^2.2.3",
"jsonwebtoken": "^9.0.1",
"module-alias": "^2.2.3",
"nsfwjs": "^3.0.0",
"pm2": "^5.3.1",
"reflect-metadata": "^0.1.13",
"semver": "^7.5.4",
"sharp": "^0.33.2",
"socket.io": "^4.7.2",
"tesseract.js": "^5.0.4",
"tslib": "^2.6.1",
"undici": "^5.23.0",
"uuid": "^9.0.0",
"zod": "^3.21.4"
},
"devDependencies": {
"@commitlint/cli": "^17.6.6",
"@commitlint/config-conventional": "^17.6.6",
"@types/bcrypt": "^5.0.0",
"@types/bun": "^1.0.5",
"@types/cors": "^2.8.13",
"@types/dot-object": "^2.1.2",
"@types/express": "^4.17.17",
"@types/figlet": "^1.5.8",
"@types/jsonwebtoken": "^9.0.2",
"@types/module-alias": "^2.0.2",
"@types/node": "^20.4.0",
"@types/semver": "^7.5.4",
"@types/uuid": "^9.0.2",
"husky": "latest",
"prisma": "^5.8.1",
"ts-node": "^10.9.1",
"typescript": "^5.3.3",
"zod-to-json-schema": "^3.21.4"
},
"optionalDependencies": {
"openai": "^4.26.0"
},
"scripts": {
"start": "node build/index.js",
"start:bun": "bun build/bun.js",
"prepare": "husky",
"dev": "bun run src/bun.ts",
"lint": "eslint src/ --ext .ts",
"lint:fix": "eslint src/ --ext .ts --fix",
"build": "tsc",
"start:prod": "pm2 start ./ecosystem.config.js",
"deploy": "node scripts/deploy-commands.js",
"gen:schema": "node scripts/generate-config-schema.js",
"clean": "rm -frv build; make clean",
"test": "node --test --require ts-node/register tests/**/*.test.ts"
},
"_moduleAliases": {
"@sudobot": "build"
}
}
78 changes: 76 additions & 2 deletions src/automod/MessageRuleService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,9 @@ import Service from "../core/Service";
import { CreateLogEmbedOptions } from "../services/LoggerService";
import { HasEventListeners } from "../types/HasEventListeners";
import { MessageRuleType } from "../types/MessageRuleSchema";
import { log, logError, logWarn } from "../utils/logger";
import { log, logDebug, logError, logWarn } from "../utils/logger";
import { escapeRegex, getEmoji, request } from "../utils/utils";
import sharp from "sharp";

export const name = "messageRuleService";

Expand All @@ -62,7 +63,8 @@ const handlers: Record<MessageRuleType["type"], RuleInfo> = {
regex_must_match: "ruleRegexMustMatch",
image: "ruleImage",
embed: "ruleEmbed",
EXPERIMENTAL_url_crawl: "ruleURLCrawl"
EXPERIMENTAL_url_crawl: "ruleURLCrawl",
EXPERIMENTAL_nsfw_filter: "ruleNSFWFilter"
};

type MessageRuleAction = MessageRuleType["actions"][number];
Expand Down Expand Up @@ -292,6 +294,78 @@ export default class MessageRuleService extends Service implements HasEventListe
return { includes: false };
}

async ruleNSFWFilter(message: Message, rule: Extract<MessageRuleType, { type: "EXPERIMENTAL_nsfw_filter" }>) {
logDebug("Scanning for NSFW content");

if (message.attachments.size === 0) {
return null;
}

const { score_thresholds } = rule;

for (const attachment of message.attachments.values()) {
logDebug("Scanning attachment", attachment.id);

if (attachment instanceof Attachment && attachment.contentType?.startsWith("image/")) {
logDebug("Scanning image attachment", attachment.id);

const [response, error] = await request({
url: attachment.proxyURL,
method: "GET",
responseType: "arraybuffer"
});

if (error || !response) {
logError(error);
return;
}

const imageData = Buffer.from(response.data, "binary");
const sharpMethodName = attachment.contentType.startsWith("image/gif")
? "gif"
: attachment.contentType.startsWith("image/png")
? "png"
: attachment.contentType.startsWith("image/webp")
? "webp"
: attachment.contentType.startsWith("image/jpeg")
? "jpeg"
: "unknown";

if (sharpMethodName === "unknown") {
logWarn("Unknown image type");
continue;
}

const sharpInfo = sharp(imageData);
const sharpMethod = sharpInfo[sharpMethodName].bind(sharpInfo);
const convertedImageBuffer = await sharpMethod().toBuffer();
const result = await this.client.imageRecognitionService.detectNSFW(convertedImageBuffer);
const isNSFW =
result.hentai >= score_thresholds.hentai ||
result.porn >= score_thresholds.porn ||
result.sexy >= score_thresholds.sexy;

logDebug("NSFW result", result);

if (isNSFW) {
return {
title: "NSFW content detected in image",
fields: [
{
name: "Scores",
value: `Hentai: ${Math.round(result.hentai * 100)}%\nPorn: ${Math.round(
result.porn * 100
)}%\nSexy: ${Math.round(result.sexy * 100)}%\nNeutral: ${Math.round(result.neutral * 100)}%`
}
]
} as CreateLogEmbedOptions;
}
}
}

return null;
}

async ruleEmbed(message: Message, rule: Extract<MessageRuleType, { type: "embed" }>) {
if (message.embeds.length === 0) {
return null;
Expand Down
37 changes: 36 additions & 1 deletion src/services/ImageRecognitionService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,37 @@

import Tesseract, { createWorker } from "tesseract.js";
import Service from "../core/Service";
import { log } from "../utils/logger";
import { log, logInfo } from "../utils/logger";
import { NSFWJS, load } from "nsfwjs";
import * as tf from "@tensorflow/tfjs-node";
import { developmentMode } from "../utils/utils";
import sharp from "sharp";

export const name = "imageRecognitionService";

if (!developmentMode()) {
tf.enableProdMode();
}

export default class ImageRecognitionService extends Service {
protected worker: Tesseract.Worker | null = null;
protected nsfwJsModel: NSFWJS | null = null;
protected timeout: Timer | null = null;

async boot() {
for (const guild in this.client.configManager.config) {
if (
this.client.configManager.config[guild]?.message_rules?.rules.some(
rule => rule.type === "EXPERIMENTAL_nsfw_filter"
)
) {
logInfo("Loading NSFWJS model for NSFW image recognition");
this.nsfwJsModel = await load(process.env.NSFWJS_MODEL_URL || undefined);
break;
}
}
}

protected async createWorkerIfNeeded() {
if (!this.worker && !this.timeout) {
log("Spawning new tesseract worker for image recognition");
Expand All @@ -51,4 +74,16 @@ export default class ImageRecognitionService extends Service {
await this.createWorkerIfNeeded();
return this.worker!.recognize(image);
}

async detectNSFW(image: Buffer) {
const tensor = tf.node.decodeImage(image, 3);
const predictions = await this.nsfwJsModel!.classify(tensor as tf.Tensor3D);
const result: Record<string, number> = {};

for (const prediction of predictions) {
result[prediction.className.toLowerCase()] = prediction.probability;
}

return result;
}
}
17 changes: 16 additions & 1 deletion src/types/MessageRuleSchema.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,20 @@ export const URLCrawlRule = z
})
.describe("Experimental. Use at your own risk.");

export const NSFWFilter = z
.object({
...Common,
type: z.literal("EXPERIMENTAL_nsfw_filter"),
score_thresholds: z
.object({
hentai: z.number().min(0).max(1).default(0.35),
porn: z.number().min(0).max(1).default(0.35),
sexy: z.number().min(0).max(1).default(0.8)
})
.default({})
})
.describe("Experimental. Use at your own risk.");

export const MessageRuleSchema = z.union([
DomainRule,
BlockedMimeTypeRule,
Expand All @@ -133,7 +147,8 @@ export const MessageRuleSchema = z.union([
RegexMustMatchRule,
ImageRule,
EmbedRule,
URLCrawlRule
URLCrawlRule,
NSFWFilter
]);

export type MessageRuleType = z.infer<typeof MessageRuleSchema>;
Loading

0 comments on commit 43eb34e

Please sign in to comment.