feat: auto nsfw detection (#8840)

* feat: auto nsfw detection

* ✌️

* Update ja-JP.yml

* Update ja-JP.yml

* ポルノ判定のしきい値を高めに

* エラーハンドリングちゃんとした

* Update ja-JP.yml

* 感度設定を強化

* refactor

* feat: add video support for auto nsfw detection

* rename: image -> media

* .js

* fix: add missing error handling

* fix: use valid pathname instead of using filename due to invalid usage

* perf(nsfw-detection): decode frames

* disable detection of video for some reasons

* perf(nsfw-detection): streamify detection process for video

* disable disallowUploadWhenPredictedAsPorn option

* fix(nsfw-detection): improve reliability

* fix(nsfw-detection): use Math.ceil instead of Math.round

* perf(nsfw-detection): delete tmp frames after used

* fix(nsfw-detection): FSWatcher does not emit ready event

* perf(nsfw-detection): skip black frames

* refactor: strip exists check

* Update package.json

* めっちゃ変えた

* lint

* Update COPYING

* オプションで動画解析できるように

* Update yarn.lock

* Update CHANGELOG.md

Co-authored-by: Acid Chicken (硫酸鶏) <root@acid-chicken.com>
This commit is contained in:
syuilo
2022-07-07 21:06:37 +09:00
committed by GitHub
parent 010db2515c
commit e560601815
39 changed files with 1275 additions and 78 deletions

View File

@ -1,12 +1,18 @@
import * as fs from 'node:fs';
import * as crypto from 'node:crypto';
import { join } from 'node:path';
import * as stream from 'node:stream';
import * as util from 'node:util';
import { FSWatcher } from 'chokidar';
import { fileTypeFromFile } from 'file-type';
import FFmpeg from 'fluent-ffmpeg';
import isSvg from 'is-svg';
import probeImageSize from 'probe-image-size';
import { type predictionType } from 'nsfwjs';
import sharp from 'sharp';
import { encode } from 'blurhash';
import { detectSensitive } from '@/services/detect-sensitive.js';
import { createTempDir } from './create-temp.js';
const pipeline = util.promisify(stream.pipeline);
@ -21,6 +27,8 @@ export type FileInfo = {
height?: number;
orientation?: number;
blurhash?: string;
sensitive: boolean;
porn: boolean;
warnings: string[];
};
@ -37,7 +45,12 @@ const TYPE_SVG = {
/**
* Get file information
*/
export async function getFileInfo(path: string): Promise<FileInfo> {
export async function getFileInfo(path: string, opts: {
skipSensitiveDetection: boolean;
sensitiveThreshold?: number;
sensitiveThresholdForPorn?: number;
enableSensitiveMediaDetectionForVideos?: boolean;
}): Promise<FileInfo> {
const warnings = [] as string[];
const size = await getFileSize(path);
@ -58,7 +71,7 @@ export async function getFileInfo(path: string): Promise<FileInfo> {
// うまく判定できない画像は octet-stream にする
if (!imageSize) {
warnings.push(`cannot detect image dimensions`);
warnings.push('cannot detect image dimensions');
type = TYPE_OCTET_STREAM;
} else if (imageSize.wUnits === 'px') {
width = imageSize.width;
@ -67,7 +80,7 @@ export async function getFileInfo(path: string): Promise<FileInfo> {
// 制限を超えている画像は octet-stream にする
if (imageSize.width > 16383 || imageSize.height > 16383) {
warnings.push(`image dimensions exceeds limits`);
warnings.push('image dimensions exceeds limits');
type = TYPE_OCTET_STREAM;
}
} else {
@ -84,6 +97,19 @@ export async function getFileInfo(path: string): Promise<FileInfo> {
});
}
let sensitive = false;
let porn = false;
if (!opts.skipSensitiveDetection) {
[sensitive, porn] = await detectSensitivity(
path,
type.mime,
opts.sensitiveThreshold ?? 0.5,
opts.sensitiveThresholdForPorn ?? 0.75,
opts.enableSensitiveMediaDetectionForVideos ?? false,
);
}
return {
size,
md5,
@ -92,10 +118,150 @@ export async function getFileInfo(path: string): Promise<FileInfo> {
height,
orientation,
blurhash,
sensitive,
porn,
warnings,
};
}
async function detectSensitivity(source: string, mime: string, sensitiveThreshold: number, sensitiveThresholdForPorn: number, analyzeVideo: boolean): Promise<[sensitive: boolean, porn: boolean]> {
let sensitive = false;
let porn = false;
function judgePrediction(result: readonly predictionType[]): [sensitive: boolean, porn: boolean] {
let sensitive = false;
let porn = false;
if ((result.find(x => x.className === 'Sexy')?.probability ?? 0) > sensitiveThreshold) sensitive = true;
if ((result.find(x => x.className === 'Hentai')?.probability ?? 0) > sensitiveThreshold) sensitive = true;
if ((result.find(x => x.className === 'Porn')?.probability ?? 0) > sensitiveThreshold) sensitive = true;
if ((result.find(x => x.className === 'Porn')?.probability ?? 0) > sensitiveThresholdForPorn) porn = true;
return [sensitive, porn];
}
if (['image/jpeg', 'image/png', 'image/webp'].includes(mime)) {
const result = await detectSensitive(source);
if (result) {
[sensitive, porn] = judgePrediction(result);
}
} else if (analyzeVideo && (mime === 'image/apng' || mime.startsWith('video/'))) {
const [outDir, disposeOutDir] = await createTempDir();
try {
const command = FFmpeg()
.input(source)
.inputOptions([
'-skip_frame', 'nokey', // 可能ならキーフレームのみを取得してほしいとする(そうなるとは限らない)
'-lowres', '3', // 元の画質でデコードする必要はないので 1/8 画質でデコードしてもよいとする(そうなるとは限らない)
])
.noAudio()
.videoFilters([
{
filter: 'select', // フレームのフィルタリング
options: {
e: 'eq(pict_type,PICT_TYPE_I)', // I-Frame のみをフィルタするVP9 とかはデコードしてみないとわからないっぽい)
},
},
{
filter: 'blackframe', // 暗いフレームの検出
options: {
amount: '0', // 暗さに関わらず全てのフレームで測定値を取る
},
},
{
filter: 'metadata',
options: {
mode: 'select', // フレーム選択モード
key: 'lavfi.blackframe.pblack', // フレームにおける暗部の百分率(前のフィルタからのメタデータを参照する)
value: '50',
function: 'less', // 50% 未満のフレームを選択する50% 以上暗部があるフレームだと誤検知を招くかもしれないので)
},
},
{
filter: 'scale',
options: {
w: 299,
h: 299,
},
},
])
.format('image2')
.output(join(outDir, '%d.png'))
.outputOptions(['-vsync', '0']); // 可変フレームレートにすることで穴埋めをさせない
const results: ReturnType<typeof judgePrediction>[] = [];
let frameIndex = 0;
let targetIndex = 0;
let nextIndex = 1;
for await (const path of asyncIterateFrames(outDir, command)) {
try {
const index = frameIndex++;
if (index !== targetIndex) {
continue;
}
targetIndex = nextIndex;
nextIndex += index; // fibonacci sequence によってフレーム数制限を掛ける
const result = await detectSensitive(path);
if (result) {
results.push(judgePrediction(result));
}
} finally {
fs.promises.unlink(path);
}
}
sensitive = results.filter(x => x[0]).length >= Math.ceil(results.length * sensitiveThreshold);
porn = results.filter(x => x[1]).length >= Math.ceil(results.length * sensitiveThresholdForPorn);
} finally {
disposeOutDir();
}
}
return [sensitive, porn];
}
async function* asyncIterateFrames(cwd: string, command: FFmpeg.FfmpegCommand): AsyncGenerator<string, void> {
const watcher = new FSWatcher({
cwd,
disableGlobbing: true,
});
let finished = false;
command.once('end', () => {
finished = true;
watcher.close();
});
command.run();
for (let i = 1; true; i++) { // eslint-disable-line @typescript-eslint/no-unnecessary-condition
const current = `${i}.png`;
const next = `${i + 1}.png`;
const framePath = join(cwd, current);
if (await exists(join(cwd, next))) {
yield framePath;
} else if (!finished) { // eslint-disable-line @typescript-eslint/no-unnecessary-condition
watcher.add(next);
await new Promise<void>((resolve, reject) => {
watcher.on('add', function onAdd(path) {
if (path === next) { // 次フレームの書き出しが始まっているなら、現在フレームの書き出しは終わっている
watcher.unwatch(current);
watcher.off('add', onAdd);
resolve();
}
});
command.once('end', resolve); // 全てのフレームを処理し終わったなら、最終フレームである現在フレームの書き出しは終わっている
command.once('error', reject);
});
yield framePath;
} else if (await exists(framePath)) {
yield framePath;
} else {
return;
}
}
}
function exists(path: string): Promise<boolean> {
return fs.promises.access(path).then(() => true, () => false);
}
/**
* Detect MIME Type and extension
*/