src/services/ImageRecognitionService.ts

/**
 * This file is part of SudoBot.
 *
 * Copyright (C) 2021-2024 OSN Developers.
 *
 * SudoBot is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * SudoBot is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with SudoBot. If not, see <https://www.gnu.org/licenses/>.
 */

import type { Tensor3D } from "@tensorflow/tfjs-node";
import type { NSFWJS } from "nsfwjs";
import Tesseract, { createWorker } from "tesseract.js";
import Service from "../core/Service";
import { log, logInfo } from "../utils/Logger";
import { developmentMode } from "../utils/utils";

export const name = "imageRecognitionService";

export default class ImageRecognitionService extends Service {
    protected worker: Tesseract.Worker | null = null;
    protected nsfwJsModel: NSFWJS | null = null;
    protected timeout: Timer | null = null;
    protected tensorFlow: typeof import("@tensorflow/tfjs-node") | null = null;
    protected nsfwJs: typeof import("nsfwjs") | 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.tensorFlow = await import("@tensorflow/tfjs-node");

                if (!developmentMode()) {
                    this.tensorFlow.enableProdMode();
                }

                this.nsfwJs = await import("nsfwjs");
                this.nsfwJsModel = await this.nsfwJs.load(
                    process.env.NSFWJS_MODEL_URL || undefined,
                    process.env.NSFWJS_MODEL_IMAGE_SIZE
                        ? {
                              size: parseInt(process.env.NSFWJS_MODEL_IMAGE_SIZE)
                          }
                        : undefined
                );

                break;
            }
        }
    }

    protected async createWorkerIfNeeded() {
        if (!this.worker && !this.timeout) {
            log("Spawning new tesseract worker for image recognition");
            this.worker = await createWorker("eng");
            this.setTimeout();
        } else if (this.worker && this.timeout) {
            log("Using existing tesseract worker for image recognition");
            clearTimeout(this.timeout);
            this.setTimeout();
        }
    }

    protected setTimeout() {
        this.timeout = setTimeout(() => {
            log("Terminating existing tesseract worker due to inactivity");
            this.worker?.terminate();
            this.timeout = null;
        }, 60_000);
    }

    async recognize(image: Tesseract.ImageLike) {
        await this.createWorkerIfNeeded();
        return this.worker!.recognize(image);
    }

    async detectNSFW(image: Uint8Array | Buffer) {
        if (!this.tensorFlow) {
            throw new Error("Tensorflow is not loaded");
        }

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

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

        tensor.dispose();
        return result;
    }
}
1	rakinar2	577	/**
2			* This file is part of SudoBot.
3			*
4			* Copyright (C) 2021-2024 OSN Developers.
5			*
6			* SudoBot is free software; you can redistribute it and/or modify it
7			* under the terms of the GNU Affero General Public License as published by
8			* the Free Software Foundation, either version 3 of the License, or
9			* (at your option) any later version.
10			*
11			* SudoBot is distributed in the hope that it will be useful, but
12			* WITHOUT ANY WARRANTY; without even the implied warranty of
13			* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14			* GNU Affero General Public License for more details.
15			*
16			* You should have received a copy of the GNU Affero General Public License
17			* along with SudoBot. If not, see <https://www.gnu.org/licenses/>.
18			*/
19
20			import type { Tensor3D } from "@tensorflow/tfjs-node";
21			import type { NSFWJS } from "nsfwjs";
22			import Tesseract, { createWorker } from "tesseract.js";
23			import Service from "../core/Service";
24			import { log, logInfo } from "../utils/Logger";
25			import { developmentMode } from "../utils/utils";
26
27			export const name = "imageRecognitionService";
28
29			export default class ImageRecognitionService extends Service {
30			protected worker: Tesseract.Worker \| null = null;
31			protected nsfwJsModel: NSFWJS \| null = null;
32			protected timeout: Timer \| null = null;
33			protected tensorFlow: typeof import("@tensorflow/tfjs-node") \| null = null;
34			protected nsfwJs: typeof import("nsfwjs") \| null = null;
35
36			async boot() {
37			for (const guild in this.client.configManager.config) {
38			if (
39			this.client.configManager.config[guild]?.message_rules?.rules.some(
40			rule => rule.type === "EXPERIMENTAL_nsfw_filter"
41			)
42			) {
43			logInfo("Loading NSFWJS model for NSFW image recognition");
44
45			this.tensorFlow = await import("@tensorflow/tfjs-node");
46
47			if (!developmentMode()) {
48			this.tensorFlow.enableProdMode();
49			}
50
51			this.nsfwJs = await import("nsfwjs");
52			this.nsfwJsModel = await this.nsfwJs.load(
53			process.env.NSFWJS_MODEL_URL \|\| undefined,
54			process.env.NSFWJS_MODEL_IMAGE_SIZE
55			? {
56			size: parseInt(process.env.NSFWJS_MODEL_IMAGE_SIZE)
57			}
58			: undefined
59			);
60
61			break;
62			}
63			}
64			}
65
66			protected async createWorkerIfNeeded() {
67			if (!this.worker && !this.timeout) {
68			log("Spawning new tesseract worker for image recognition");
69			this.worker = await createWorker("eng");
70			this.setTimeout();
71			} else if (this.worker && this.timeout) {
72			log("Using existing tesseract worker for image recognition");
73			clearTimeout(this.timeout);
74			this.setTimeout();
75			}
76			}
77
78			protected setTimeout() {
79			this.timeout = setTimeout(() => {
80			log("Terminating existing tesseract worker due to inactivity");
81			this.worker?.terminate();
82			this.timeout = null;
83			}, 60_000);
84			}
85
86			async recognize(image: Tesseract.ImageLike) {
87			await this.createWorkerIfNeeded();
88			return this.worker!.recognize(image);
89			}
90
91			async detectNSFW(image: Uint8Array \| Buffer) {
92			if (!this.tensorFlow) {
93			throw new Error("Tensorflow is not loaded");
94			}
95
96			const tensor = this.tensorFlow.node.decodeImage(image, 3, undefined, false);
97			const predictions = await this.nsfwJsModel!.classify(tensor as Tensor3D);
98			const result: Record<string, number> = {};
99
100			for (const prediction of predictions) {
101			result[prediction.className.toLowerCase()] = prediction.probability;
102			}
103
104			tensor.dispose();
105			return result;
106			}
107			}