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 { + public async detectSensitive(source: string | Buffer): Promise { 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];