Add handling of NSFW skins

packages/skin-database/addSkin.ts

@@ -6,6 +6,7 @@ import Shooter from "./shooter";
 import _temp from "temp";
 import * as Analyser from "./analyser";
 
+// TODO Move this into the function so that we clean up on each run?
 const temp = _temp.track();
 
 // TODO

packages/skin-database/cli.ts

@@ -12,6 +12,8 @@ import db from "./db";
 import Discord from "discord.js";
 import { tweet } from "./tasks/tweet";
 import { addSkinFromBuffer } from "./addSkin";
+import fetch from "node-fetch";
+import { analyseBuffer, NsfwPrediction } from "./nsfwImage";
 
 async function main() {
   const client = new Discord.Client();
@@ -69,9 +71,47 @@ async function main() {
         console.log(await addSkinFromBuffer(buffer, filePath, "cli-user"));
         break;
       }
+      case "nsfw": {
+        console.log(await Skins.getSkinToReviewForNsfw());
+        break;
+      }
+      case "confirm-nsfw-predictions": {
+        const md5s = await Skins.getMissingNsfwPredictions();
+        console.log(`Found ${md5s.length} to predict`);
+
+        for (const md5 of md5s) {
+          try {
+            await Skins.computeAndSetNsfwPredictions(md5);
+          } catch (e) {
+            console.error(e);
+          }
+        }
+        console.log("Done.");
+        break;
+      }
       case "migrate": {
         await migrate();
       }
+      case "tweet-data": {
+        // From running `tweet.py sort`
+        const file = fs.readFileSync(
+          path.join(__dirname, "../../tweetBot/likes.txt"),
+          { encoding: "utf8" }
+        );
+
+        const lines = file.split("\n");
+        for (const line of lines) {
+          if (line == null || line === "") {
+            return;
+          }
+          const [md5, likes, tweetId] = line.split(" ");
+          console.log({ md5, likes, tweetId });
+          await Skins.setTweetInfo(md5, Number(likes), tweetId);
+        }
+
+        console.log("done");
+        break;
+      }
 
       default:
         console.log(`Unknown command ${argv._[0]}`);

packages/skin-database/data/skins.ts

@@ -3,6 +3,8 @@ import path from "path";
 import S3 from "../s3";
 import logger from "../logger";
 import { DBSkinRecord, SkinRecord, DBIARecord, TweetStatus } from "../types";
+import fetch from "node-fetch";
+import { analyseBuffer, NsfwPrediction } from "../nsfwImage";
 
 const skins = db.get("skins");
 const iaItems = db.get("internetArchiveItems");
@@ -39,6 +41,7 @@ function getSkinRecord(skin: DBSkinRecord): SkinRecord {
     rejected,
     approved,
     nsfw,
+    nsfwPredictions,
   } = skin;
   const fileNames = filePaths.map((p) => path.basename(p));
   const skinUrl = `https://s3.amazonaws.com/webamp-uploaded-skins/skins/${md5}.wsz`;
@@ -60,6 +63,7 @@ function getSkinRecord(skin: DBSkinRecord): SkinRecord {
     rejected,
     approved,
     nsfw,
+    nsfwPredictions,
   };
 }
 
@@ -175,6 +179,14 @@ export async function getUnarchived() {
   return skins.find({ itemName: null }, { md5: 1 });
 }
 
+export async function getMissingNsfwPredictions() {
+  const results = await skins.find(
+    { nsfwPredictions: null, type: "CLASSIC" },
+    { md5: 1 }
+  );
+  return results.map(({ md5 }) => md5);
+}
+
 export function getInternetArchiveUrl(itemName: string | null): string | null {
   return itemName == null ? null : `https://archive.org/details/${itemName}`;
 }
@@ -254,6 +266,19 @@ export async function getSkinToReview(): Promise<{
   return { filename: canonicalFilename, md5 };
 }
 
+export async function getSkinToReviewForNsfw(): Promise<{
+  filename: string | null;
+  md5: string;
+}> {
+  const reviewable = await skins.find(REVIEWABLE_QUERY, {
+    limit: 1,
+    sort: { "nsfwPredictions.porn": -1 },
+  });
+  const skin = reviewable[0];
+  const { canonicalFilename, md5 } = getSkinRecord(skin);
+  return { filename: canonicalFilename, md5 };
+}
+
 export async function getSkinToTweet(): Promise<SkinRecord | null> {
   const tweetables = await skins.aggregate([
     { $match: TWEETABLE_QUERY },
@@ -315,3 +340,71 @@ export async function getRandomClassicSkinMd5() {
   }
   return random[0].md5;
 }
+
+export async function getScreenshotBuffer(md5: string): Promise<Buffer> {
+  const skin = await getSkinByMd5(md5);
+  if (skin == null) {
+    throw new Error(`Could not find skin with hash ${md5}`);
+  }
+  const screenshotResponse = await fetch(skin?.screenshotUrl);
+  if (!screenshotResponse.ok) {
+    throw new Error(`Could not get screenshot at ${skin?.screenshotUrl}`);
+  }
+  return screenshotResponse.buffer();
+}
+
+export async function setNsfwPredictions(
+  md5: string,
+  nsfwPredictions: NsfwPrediction
+): Promise<void> {
+  await skins.findOneAndUpdate({ md5 }, { $set: { nsfwPredictions } });
+}
+
+export async function setTweetInfo(
+  md5: string,
+  likes: number,
+  tweetId: string
+): Promise<void> {
+  if (md5 === "48bbdbbeb03d347e59b1eebda4d352d0") {
+    console.log(likes, tweetId);
+  }
+  await skins.findOneAndUpdate(
+    { md5 },
+    { $set: { twitterLikes: likes, tweetId } }
+  );
+}
+
+export async function computeAndSetNsfwPredictions(md5: string): Promise<void> {
+  const image = await getScreenshotBuffer(md5);
+  const predictions = await analyseBuffer(image);
+  await setNsfwPredictions(md5, predictions);
+}
+
+export async function getMuseumPage({
+  offset,
+  first,
+}: {
+  offset: number;
+  first: number;
+}): Promise<
+  Array<{ color: string; fileName: string; md5: string; nsfw: boolean }>
+> {
+  const reviewable = await skins.find(
+    { type: "CLASSIC" },
+    {
+      limit: first,
+      skip: offset,
+      sort: { twitterLikes: -1, approved: -1, rejected: 1 },
+      fields: { averageColor: 1, md5: 1, nsfw: 1 },
+    }
+  );
+
+  return reviewable.map(({ md5, averageColor, nsfw }) => {
+    return {
+      color: averageColor,
+      filename: "FILENAME",
+      md5,
+      nsfw,
+    };
+  });
+}

packages/skin-database/discord-bot/commands/review.ts

@@ -2,8 +2,10 @@ import * as Skins from "../../data/skins";
 import * as Utils from "../utils";
 import { Message } from "discord.js";
 
-async function reviewSkin(message: Message): Promise<void> {
-  const skin = await Skins.getSkinToReview();
+async function reviewSkin(message: Message, nsfw: boolean): Promise<void> {
+  const skin = await (nsfw
+    ? Skins.getSkinToReviewForNsfw()
+    : Skins.getSkinToReview());
   if (skin == null) {
     throw new Error("No skins to review");
   }
@@ -15,8 +17,9 @@ async function reviewSkin(message: Message): Promise<void> {
   });
 }
 
-async function handler(message: Message, args: [string]) {
+async function handler(message: Message, args: [string, string]) {
   let count = Number(args[0] || 1);
+  let nsfw = args[1] === "nsfw";
   if (count > 50) {
     await message.channel.send(
       `You can only review up to ${count} skins at a time.`
@@ -26,7 +29,7 @@ async function handler(message: Message, args: [string]) {
   await message.channel.send(`Going to show ${count} skins to review.`);
   let i = count;
   while (i--) {
-    await reviewSkin(message);
+    await reviewSkin(message, nsfw);
   }
   if (count > 1) {
     const tweetableCount = await Skins.getTweetableSkinCount();

packages/skin-database/index.js

@@ -1,5 +1,6 @@
 const express = require("express");
 const app = express();
+const config = require("./config");
 const db = require("./db");
 const iaItems = db.get("internetArchiveItems");
 // const info = require("/Volumes/Mobile Backup/skins/cache/info.json");
@@ -8,6 +9,23 @@ const port = 3001;
 const graphql = require("./graphql").default;
 const fileUpload = require("express-fileupload");
 const { addSkinFromBuffer } = require("./addSkin");
+const Discord = require("discord.js");
+const Utils = require("./discord-bot/utils");
+const cors = require("cors");
+
+const whitelist = ["https://skins.webamp.org", "http://localhost:3000"];
+const corsOptions = {
+  origin: function (origin, callback) {
+    if (whitelist.indexOf(origin) !== -1 || !origin) {
+      console.log("okay");
+      callback(null, true);
+    } else {
+      callback(new Error("Not allowed by CORS"));
+    }
+  },
+};
+
+app.use(cors(corsOptions));
 
 // TODO: Look into 766c4fad9088037ab4839b18292be8b1
 // Has huge number of filenames in info.json
@@ -33,6 +51,18 @@ app.get("/items/:identifier", async (req, res) => {
   res.json(item);
 });
 
+app.get("/skins/", async (req, res) => {
+  const { offset = 0, first = 100 } = req.query;
+  const [skins, skinCount] = await Promise.all([
+    Skins.getMuseumPage({
+      offset: Number(offset),
+      first: Number(first),
+    }),
+    Skins.getClassicSkinCount(),
+  ]);
+  res.json({ skinCount, skins });
+});
+
 app.post("/skins/", async (req, res) => {
   const files = req.files;
   if (files == null) {
@@ -58,6 +88,22 @@ app.get("/skins/:md5", async (req, res) => {
   res.json(skin);
 });
 
+// TODO: Make this POST
+app.post("/skins/:md5/report", async (req, res) => {
+  const { md5 } = req.params;
+  const client = new Discord.Client();
+  await client.login(config.discordToken);
+  const dest = client.channels.get(config.NSFW_SKIN_CHANNEL_ID);
+
+  // Don't await
+  Utils.postSkin({
+    md5,
+    title: (filename) => `Review: ${filename}`,
+    dest,
+  });
+  res.send("The skin has been reported and will be reviewed shortly.");
+});
+
 app.get("/skins/:md5/readme.txt", async (req, res) => {
   const { md5 } = req.params;
   const readmeText = await Skins.getReadme(md5);

packages/skin-database/nsfw.js

@@ -2,20 +2,11 @@ const tf = require("@tensorflow/tfjs-node");
 const nsfw = require("nsfwjs");
 const fs = require("fs");
 const path = require("path");
-const files = [
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/b905cb364791fe989e811492e84c5c2a.png",
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/6d77accd3add234e8e4dca93fc102ea8.png",
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/820beeda9c199fdac42e4071de8b3331.png",
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/4d1248a2629a17591008cbb9f1e7ac20.png",
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/3e83cd504a101c513ea9cf8b0606fd3d.png",
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/746ddb3108f11a80f96b4cd9e27e53ee.png",
-  "/Users/jordaneldredge/Downloads/nsfw-detection-test/71124781f08ff3d4069de62450c90283.png",
-];
 
-const skinsDir = "/Volumes/Mobile Backup/skins/md5Screenshots";
+const modelPromise = nsfw.load();
 
 async function fn() {
-  const model = await nsfw.load(); // To load a local model, nsfw.load('file://./path/to/model/')
+  const model = await modelPromise; // To load a local model, nsfw.load('file://./path/to/model/')
   const screenshots = fs.readdirSync(skinsDir);
   for (const file of screenshots) {
     if (!file.endsWith(".png")) {

packages/skin-database/nsfwImage.ts

@@ -0,0 +1,30 @@
+import * as tf from "@tensorflow/tfjs-node";
+import { load as nsfwLoad } from "nsfwjs";
+import path from "path";
+import fs from "fs";
+
+const modelPromise = nsfwLoad();
+
+export type NsfwPrediction = {
+  porn: number;
+  neutral: number;
+  sexy: number;
+  hentai: number;
+  drawing: number;
+};
+
+export async function analyseBuffer(buffer: Buffer): Promise<NsfwPrediction> {
+  const model = await modelPromise;
+  const image = await tf.node.decodePng(buffer, 3);
+  const predictions = await model.classify(image);
+  if (predictions.length < 5) {
+    throw new Error("Not enough predictions found");
+  }
+
+  // @ts-ignore
+  const obj: NsfwPrediction = {};
+  predictions.forEach((prediction) => {
+    obj[prediction.className.toLowerCase()] = prediction.probability;
+  });
+  return obj;
+}

packages/skin-database/package.json

@@ -5,7 +5,9 @@
   "license": "MIT",
   "dependencies": {
     "@tensorflow/tfjs-node": "^2.0.0",
+    "algoliasearch": "^4.3.0",
     "aws-sdk": "^2.663.0",
+    "cors": "^2.8.5",
     "discord.js": "^11.4.2",
     "express": "^4.17.1",
     "express-fileupload": "^1.1.7-alpha.3",
@@ -30,7 +32,8 @@
     "tweet": "tsc && ./dist/cli.js tweet",
     "fetch-metadata": "tsc && .dist/cli.js fetch-metadata",
     "bot": "tsc && node ./dist/discord-bot/index.js",
-    "cli": "tsc && node ./dist/cli.js"
+    "cli": "tsc && node ./dist/cli.js",
+    "index": "tsc && node ./dist/tasks/searchIndex.js"
   },
   "prettier": {},
   "devDependencies": {

packages/skin-database/tasks/searchIndex.js

@@ -0,0 +1,76 @@
+const algoliasearch = require("algoliasearch");
+const { ALGOLIA_INDEX, ALGOLIA_KEY } = require("../config");
+const path = require("path");
+const db = require("../db");
+
+const client = algoliasearch(ALGOLIA_INDEX, ALGOLIA_KEY);
+const index = client.initIndex("Skins");
+
+function tuncate(str, len) {
+  const overflow = str.length - len;
+  if (overflow < 0) {
+    return str;
+  }
+
+  const half = Math.floor((len - 1) / 2);
+
+  const start = str.slice(0, half);
+  const end = str.slice(-half);
+  return `${start} ########### ${end}`;
+}
+
+function buildSkinIndex(skin) {
+  const { filePaths, nsfw } = skin;
+  if (!filePaths || filePaths.length === 0) {
+    console.warn("no file name for ", md5);
+    return;
+  }
+  const fileName = path.basename(filePaths[0]);
+  const readmeText = skin.readmeText ? tuncate(skin.readmeText, 4800) : null;
+  return {
+    objectID: skin.md5,
+    nsfw,
+    //md5,
+    //fileName,
+    // emails: skin.emails || null,
+    // readmeText
+    // color: skin.averageColor
+    // twitterLikes: Number(skin.twitterLikes || 0)
+  };
+}
+
+async function go({ dry = true }) {
+  const skins = await db.get("skins").find(
+    { type: "CLASSIC" },
+    {
+      fields: {
+        md5: 1,
+        averageColor: 1,
+        nsfw: 1,
+        twitterLikes: 1,
+        readmeText: 1,
+        filePaths: 1,
+      },
+    }
+  );
+
+  const indexes = skins.map(buildSkinIndex).filter(Boolean);
+  // .filter((index) => index.nsfw);
+
+  db.close();
+
+  if (dry) {
+    console.log("Index turned off. Turn it on if you really mean it");
+    return;
+  }
+  console.log("Writing index");
+  const results = await new Promise((resolve, reject) => {
+    index.partialUpdateObjects(indexes, function (err, content) {
+      if (err != null) reject(err);
+      resolve(content);
+    });
+  });
+  console.log("done!", results);
+}
+
+go({ dry: true }); // .then(content => console.log("Updated index for:", content.length));

packages/skin-database/types.ts

@@ -1,3 +1,5 @@
+import { NsfwPrediction } from "./nsfwImage";
+
 export type TweetStatus = "APPROVED" | "REJECTED" | "TWEETED" | "UNREVIEWED";
 
 export type DBSkinRecord = {
@@ -14,6 +16,7 @@ export type DBSkinRecord = {
   rejected?: boolean;
   approved?: boolean;
   nsfw?: boolean;
+  nsfwPredictions?: NsfwPrediction;
 };
 
 export type DBIARecord = {
@@ -38,4 +41,5 @@ export type SkinRecord = {
   rejected?: boolean;
   approved?: boolean;
   nsfw?: boolean;
+  nsfwPredictions?: NsfwPrediction;
 };