blob: 248a9b8979a624d32cd1fc953ceb3c693e0412da (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
|
/*
* SPDX-FileCopyrightText: syuilo and misskey-project
* SPDX-License-Identifier: AGPL-3.0-only
*/
import * as fs from 'node:fs';
import { fileURLToPath } from 'node:url';
import { dirname } from 'node:path';
import { Injectable } from '@nestjs/common';
import * as nsfw from 'nsfwjs';
import si from 'systeminformation';
import { Mutex } from 'async-mutex';
import fetch from 'node-fetch';
import { bindThis } from '@/decorators.js';
const _filename = fileURLToPath(import.meta.url);
const _dirname = dirname(_filename);
const REQUIRED_CPU_FLAGS_X64 = ['avx2', 'fma'];
let isSupportedCpu: undefined | boolean = undefined;
@Injectable()
export class AiService {
private model: nsfw.NSFWJS;
private modelLoadMutex: Mutex = new Mutex();
constructor(
) {
}
@bindThis
public async detectSensitive(path: string): Promise<nsfw.PredictionType[] | null> {
try {
if (isSupportedCpu === undefined) {
isSupportedCpu = await this.computeIsSupportedCpu();
}
if (!isSupportedCpu) {
console.error('This CPU cannot use TensorFlow.');
return null;
}
const tf = await import('@tensorflow/tfjs-node');
tf.env().global.fetch = fetch;
if (this.model == null) {
await this.modelLoadMutex.runExclusive(async () => {
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;
}
}
private async computeIsSupportedCpu(): Promise<boolean> {
switch (process.arch) {
case 'x64': {
const cpuFlags = await this.getCpuFlags();
return REQUIRED_CPU_FLAGS_X64.every(required => cpuFlags.includes(required));
}
case 'arm64': {
// As far as I know, no required CPU flags for ARM64.
return true;
}
default: {
return false;
}
}
}
@bindThis
private async getCpuFlags(): Promise<string[]> {
const str = await si.cpuFlags();
return str.split(/\s+/);
}
}
|
