fix(backend): webpなどの画像に対してセンシティブなメディアの検出が適用されていなかった問題を修正 (#16523)

画像をnsfwjsにかける前にsharpで均一にするようにした

2 files changed, 21 insertions, 12 deletions

diff --git a/packages/backend/src/core/AiService.ts b/packages/backend/src/core/AiService.ts
index 248a9b8979..23ab8082ed 100644
--- a/packages/backend/src/core/AiService.ts
+++ b/packages/backend/src/core/AiService.ts
@@ -29,7 +29,7 @@ export class AiService {
 	}
 
 	@bindThis
-	public async detectSensitive(path: string): Promise<nsfw.PredictionType[] | null> {
+	public async detectSensitive(source: string | Buffer): Promise<nsfw.PredictionType[] | null> {
 		try {
 			if (isSupportedCpu === undefined) {
 				isSupportedCpu = await this.computeIsSupportedCpu();
@@ -51,7 +51,7 @@ export class AiService {
 				});
 			}
 
-			const buffer = await fs.promises.readFile(path);
+			const buffer = source instanceof Buffer ? source : await fs.promises.readFile(source);
 			const image = await tf.node.decodeImage(buffer, 3) as any;
 			try {
 				const predictions = await this.model.classify(image);
diff --git a/packages/backend/src/core/FileInfoService.ts b/packages/backend/src/core/FileInfoService.ts
index 6250d4d3a1..62a7d24afb 100644
--- a/packages/backend/src/core/FileInfoService.ts
+++ b/packages/backend/src/core/FileInfoService.ts
@@ -21,6 +21,7 @@ import { LoggerService } from '@/core/LoggerService.js';
 import type Logger from '@/logger.js';
 import { bindThis } from '@/decorators.js';
 import type { PredictionType } from 'nsfwjs';
+import { isMimeImage } from '@/misc/is-mime-image.js';
 
 export type FileInfo = {
 	size: number;
@@ -204,16 +205,7 @@ export class FileInfoService {
 			return [sensitive, porn];
 		}
 
-		if ([
-			'image/jpeg',
-			'image/png',
-			'image/webp',
-		].includes(mime)) {
-			const result = await this.aiService.detectSensitive(source);
-			if (result) {
-				[sensitive, porn] = judgePrediction(result);
-			}
-		} else if (analyzeVideo && (mime === 'image/apng' || mime.startsWith('video/'))) {
+		if (analyzeVideo && (mime === 'image/apng' || mime.startsWith('video/'))) {
 			const [outDir, disposeOutDir] = await createTempDir();
 			try {
 				const command = FFmpeg()
@@ -281,6 +273,23 @@ export class FileInfoService {
 			} finally {
 				disposeOutDir();
 			}
+		} else if (isMimeImage(mime, 'sharp-convertible-image-with-bmp')) {
+			/*
+			 * tfjs-node は限られた画像形式しか受け付けないため、sharp で PNG に変換する
+			 * せっかくなので内部処理で使われる最大サイズの299x299に事前にリサイズする
+			 */
+			const png = await (await sharpBmp(source, mime))
+				.resize(299, 299, {
+					withoutEnlargement: false,
+				})
+				.rotate()
+				.flatten({ background: { r: 119, g: 119, b: 119 } }) // 透過部分を18%グレーで塗りつぶす
+				.png()
+				.toBuffer();
+			const result = await this.aiService.detectSensitive(png);
+			if (result) {
+				[sensitive, porn] = judgePrediction(result);
+			}
 		}
 
 		return [sensitive, porn];