Sharkey/packages/backend/src/core/AiService.ts

77 lines
1.9 KiB
TypeScript
Raw Normal View History

/*
* SPDX-FileCopyrightText: syuilo and other misskey contributors
* SPDX-License-Identifier: AGPL-3.0-only
*/
2022-09-17 21:27:08 +03:00
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 { Mutex } from 'async-mutex';
2022-09-20 23:33:11 +03:00
import type { Config } from '@/config.js';
2022-09-17 21:27:08 +03:00
import { DI } from '@/di-symbols.js';
2022-12-04 10:05:32 +02:00
import { bindThis } from '@/decorators.js';
2022-09-17 21:27:08 +03:00
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 {
2022-09-18 21:11:50 +03:00
private model: nsfw.NSFWJS;
private modelLoadMutex: Mutex = new Mutex();
2022-09-17 21:27:08 +03:00
constructor(
@Inject(DI.config)
private config: Config,
) {
}
@bindThis
2022-09-17 21:27:08 +03:00
public async detectSensitive(path: string): Promise<nsfw.predictionType[] | null> {
try {
if (isSupportedCpu === undefined) {
2022-09-18 21:11:50 +03:00
const cpuFlags = await this.getCpuFlags();
2022-09-17 21:27:08 +03:00
isSupportedCpu = REQUIRED_CPU_FLAGS.every(required => cpuFlags.includes(required));
}
2022-09-17 21:27:08 +03:00
if (!isSupportedCpu) {
console.error('This CPU cannot use TensorFlow.');
return null;
}
2022-09-17 21:27:08 +03:00
const tf = await import('@tensorflow/tfjs-node');
if (this.model == null) {
await this.modelLoadMutex.runExclusive(async () => {
if (this.model == null) {
this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 });
}
});
}
2022-09-17 21:27:08 +03:00
const buffer = await fs.promises.readFile(path);
const image = await tf.node.decodeImage(buffer, 3) as any;
try {
2022-09-18 21:11:50 +03:00
const predictions = await this.model.classify(image);
2022-09-17 21:27:08 +03:00
return predictions;
} finally {
image.dispose();
}
} catch (err) {
console.error(err);
return null;
}
}
@bindThis
2022-09-18 21:11:50 +03:00
private async getCpuFlags(): Promise<string[]> {
2022-09-17 21:27:08 +03:00
const str = await si.cpuFlags();
return str.split(/\s+/);
}
}