Upload 4 files

src/NsfwImageClassifier.ts

@@ -0,0 +1,46 @@
+import * as tf from '@tensorflow/tfjs-node';
+import * as nsfwjs from 'nsfwjs';
+import {NSFWJS} from 'nsfwjs';
+import {Tensor3D} from '@tensorflow/tfjs';
+import {model as config} from 'app/config/model';
+
+tf.enableProdMode();
+
+export class NsfwImageClassifier {
+  #model?: NSFWJS;
+
+  async classify(imageBuffer: Buffer) {
+    const [model, image] = await Promise.all([
+      this.#getModel(),
+      tf.node.decodeImage(imageBuffer, 3),
+    ]);
+
+    const predictions = await model.classify(image as Tensor3D);
+
+    image.dispose();
+
+    return this.#transformData(predictions);
+  }
+
+  async classifyMany(imagesBuffers: Buffer[]) {
+    return await Promise.all(imagesBuffers.map(buffer => this.classify(buffer)));
+  }
+
+  async #getModel(): Promise<NSFWJS> {
+    if (!this.#model) {
+      this.#model = await nsfwjs.load('file://model/', {size: config.size});
+    }
+
+    return this.#model;
+  }
+
+  #transformData(data: { className: string; probability: number }[]): Record<string, number> {
+    const result: Record<string, number> = {};
+
+    for (const item of data) {
+      result[item.className.toLowerCase()] = item.probability;
+    }
+
+    return result;
+  }
+}

src/config/model.ts

@@ -0,0 +1,5 @@
+const DEFAULT_MODEL_SIZE = 299;
+
+export const model = {
+  size: Number(process.env.MODEL_SIZE) || DEFAULT_MODEL_SIZE,
+};

src/controllers/NsfwController.ts

@@ -0,0 +1,42 @@
+import {Controller, Post} from 'simple-ts-express-decorators';
+import multer, {memoryStorage} from 'multer';
+import {Request, Response} from 'express';
+import {NsfwImageClassifier} from 'app/NsfwImageClassifier';
+
+const upload = multer({storage: memoryStorage()});
+
+@Controller()
+export class NsfwController {
+  classifier: NsfwImageClassifier;
+
+  constructor() {
+    this.classifier = new NsfwImageClassifier();
+  }
+
+  @Post('/classify', upload.single('image'))
+  async classify(request: Request, response: Response) {
+    if (!request.file) {
+      return response
+        .status(410)
+        .json({error: 'Specify image'});
+    }
+
+    const data = await this.classifier.classify(request.file.buffer);
+
+    return response.json(data);
+  }
+
+  @Post('/classify-many', upload.array('images', 10))
+  async classifyMany(request: Request, response: Response) {
+    if (!request.files || !request.files.length) {
+      return response
+        .status(410)
+        .json({error: 'Specify images'});
+    }
+
+    const buffers = (request.files as Express.Multer.File[]).map(file => file.buffer);
+    const data = await this.classifier.classifyMany(buffers);
+
+    return response.json(data);
+  }
+}

src/index.ts

@@ -0,0 +1,17 @@
+import 'reflect-metadata';
+import 'module-alias/register';
+
+import express from 'express';
+import bodyParser from 'body-parser';
+import {ControllersLoader} from 'simple-ts-express-decorators';
+import {NsfwController} from 'app/controllers/NsfwController';
+
+const app = express();
+
+app.use(bodyParser.json());
+
+new ControllersLoader({
+  controllers: [NsfwController]
+}).load(app);
+
+app.listen(3000);