From b75184ec8e3436200bacdcd832e3324702553d20 Mon Sep 17 00:00:00 2001
From: syuilo <Syuilotan@yahoo.co.jp>
Date: Sun, 18 Sep 2022 03:27:08 +0900
Subject: なんかもうめっちゃ変えた
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 packages/backend/src/core/AiService.ts | 60 ++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 packages/backend/src/core/AiService.ts

(limited to 'packages/backend/src/core/AiService.ts')

diff --git a/packages/backend/src/core/AiService.ts b/packages/backend/src/core/AiService.ts
new file mode 100644
index 0000000000..1cfc3382a9
--- /dev/null
+++ b/packages/backend/src/core/AiService.ts
@@ -0,0 +1,60 @@
+import * as fs from 'node:fs';
+import { fileURLToPath } from 'node:url';
+import { dirname } from 'node:path';
+import { Inject, Injectable } from '@nestjs/common';
+import * as nsfw from 'nsfwjs';
+import si from 'systeminformation';
+import { Config } from '@/config.js';
+import { DI } from '@/di-symbols.js';
+
+const _filename = fileURLToPath(import.meta.url);
+const _dirname = dirname(_filename);
+
+const REQUIRED_CPU_FLAGS = ['avx2', 'fma'];
+let isSupportedCpu: undefined | boolean = undefined;
+
+@Injectable()
+export class AiService {
+	#model: nsfw.NSFWJS;
+
+	constructor(
+		@Inject(DI.config)
+		private config: Config,
+	) {
+	}
+
+	public async detectSensitive(path: string): Promise<nsfw.predictionType[] | null> {
+		try {
+			if (isSupportedCpu === undefined) {
+				const cpuFlags = await this.#getCpuFlags();
+				isSupportedCpu = REQUIRED_CPU_FLAGS.every(required => cpuFlags.includes(required));
+			}
+	
+			if (!isSupportedCpu) {
+				console.error('This CPU cannot use TensorFlow.');
+				return null;
+			}
+	
+			const tf = await import('@tensorflow/tfjs-node');
+	
+			if (this.#model == null) this.#model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 });
+	
+			const buffer = await fs.promises.readFile(path);
+			const image = await tf.node.decodeImage(buffer, 3) as any;
+			try {
+				const predictions = await this.#model.classify(image);
+				return predictions;
+			} finally {
+				image.dispose();
+			}
+		} catch (err) {
+			console.error(err);
+			return null;
+		}
+	}
+
+	async #getCpuFlags(): Promise<string[]> {
+		const str = await si.cpuFlags();
+		return str.split(/\s+/);
+	}
+}
-- 
cgit v1.2.3-freya