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	/**
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	}