瀏覽代碼

Add manual face morphing

denikson 6 年之前
父節點
當前提交
b07180dca4
共有 5 個文件被更改,包括 127 次插入13 次删除
  1. 106 13
      commands/facemorph.js
  2. 3 0
      commands/guide.js
  3. 6 0
      commands/image_statistics.js
  4. 11 0
      db.js
  5. 1 0
      package.json

+ 106 - 13
commands/facemorph.js

@@ -4,36 +4,99 @@ const db = require("../db.js");
 const util = require("../util.js");
 const Jimp = require("jimp");
 const client = require("../client.js");
-
+const schedule = require("node-schedule");
 
 const EMOTE_GUILD = "505333548694241281";
 const clarifaiApp = new Clarifai.App({
     apiKey: ClarifaiTOKEN
 });
+const MAX_REQUESTS_PER_MONTH = 5000;
+
+let quotaRecomputeJob = null;
+let dailyResetJob = null;
+
+function recomputeQuotas() {
+    console.log(`Recomputing quotas on ${new Date().toISOString()}!`);
+    db.set("faceEditInfo.quotaRecomputeDate", "").write();
+
+    let channelPredictProbability = db.get("faceEditChannels").value();
+    let channelStatistics = db.get("faceEditStatistics").value();
+
+    let totalRandomPredicts = 0;
+
+    for (const channel in channelStatistics) {
+        if (channelStatistics.hasOwnProperty(channel) && channelPredictProbability.hasOwnProperty(channel)) {
+            totalRandomPredicts += channelPredictProbability[channel] * channelStatistics[channel];
+            channelStatistics[channel] = 0;
+        }
+    }
+
+    let dailyQuota = Math.floor((MAX_REQUESTS_PER_MONTH - totalRandomPredicts) / 30.0);
+
+    db.set("faceEditInfo.customRequestsPerDay", dailyQuota).write();
+    db.set("faceEditStatistics", channelStatistics).write();
+
+    quotaRecomputeJob.cancel();
+}
+
+function resetDailyQuota() {
+    console.log("Resetting today image request count!");
+    db.set("faceEditInfo.todayRequests", 0).write();
+}
+
+function setupJobs() {
+    let recomputeJobDate = db.get("faceEditInfo.quotaRecomputeDate").value();
+    if (recomputeJobDate)
+        quotaRecomputeJob = schedule.scheduleJob(new Date(recomputeJobDate), recomputeQuotas);
+
+    let rule = new schedule.RecurrenceRule();
+    rule.hour = 0;
+    rule.minute = 0;
+    dailyResetJob = schedule.scheduleJob(rule, resetDailyQuota);
+
+    checkQuotaJobs();
+}
+setupJobs();
+
+function checkQuotaJobs() {
+    let jobDateText = db.get("faceEditInfo.quotaRecomputeDate").value();
+    if (jobDateText)
+        return;
+
+    let jobDate = new Date();
+    jobDate.setMonth(jobDate.getMonth() + 1);
+    db.set("faceEditInfo.quotaRecomputeDate", jobDate.toISOString()).write();
+    quotaRecomputeJob.reschedule(jobDate, recomputeQuotas);
+}
 
-async function processFaceSwap(message, attachmentUrl) {
+async function processFaceSwap(message, attachmentUrl, failMessage, successMessage) {
     let result = await clarifaiApp.models.predict(Clarifai.FACE_DETECT_MODEL, attachmentUrl);
 
-    if(result.outputs[0].data.regions === undefined || result.outputs[0].data.regions.length == 0)
+    checkQuotaJobs();
+
+    if (result.outputs[0].data.regions === undefined || result.outputs[0].data.regions.length == 0) {
+        if (failMessage)
+            message.channel.send(failMessage);
         return;
-    
+    }
+
     let image = await Jimp.read(attachmentUrl);
     let w = image.getWidth();
     let h = image.getHeight();
 
     let emojiKeys = [...client.guilds.get(EMOTE_GUILD).emojis.filter(e => !e.animated).keys()];
 
-    for(let region of result.outputs[0].data.regions) {
+    for (let region of result.outputs[0].data.regions) {
         let bb = region.region_info.bounding_box;
         let bw = (bb.right_col - bb.left_col) * w;
         let bh = (bb.bottom_row - bb.top_row) * h;
-        
+
         let dx = (bb.right_col + bb.left_col) * w / 2;
         let dy = (bb.bottom_row + bb.top_row) * h / 2;
 
         let emojiKey = emojiKeys[Math.floor(Math.random() * emojiKeys.length)];
         let emoji = client.emojis.get(emojiKey);
-        
+
         let emojiImage = await Jimp.read(emoji.url);
         let ew = emojiImage.getWidth();
         let eh = emojiImage.getHeight();
@@ -50,8 +113,10 @@ async function processFaceSwap(message, attachmentUrl) {
     image.quality(90);
     let buffer = await image.getBufferAsync(Jimp.MIME_JPEG);
 
-    message.channel.send(`I noticed a face in the image. I think this looks better ${client.emojis.get("505076258753740810").toString()}`, {
-        files: [ buffer ]
+    let messageContents = successMessage || `I noticed a face in the image. I think this looks better ${client.emojis.get("505076258753740810").toString()}`;
+
+    message.channel.send(messageContents, {
+        files: [buffer]
     });
 }
 
@@ -59,9 +124,9 @@ const onMessage = (msg, contents, actionsDone) => {
     if (actionsDone)
         return false;
 
-    let imagesCount = msg.attachments.filter(v => util.isValidImage(v.filename)).size;
+    let imageAttachment = msg.attachments.find(v => util.isValidImage(v.filename));
 
-    if (imagesCount > 0) {
+    if (imageAttachment) {
         let probValue = db.get("faceEditChannels").get(msg.channel.id);
         if (probValue.isUndefined().value() || probValue.isNull().value())
             return false;
@@ -69,7 +134,6 @@ const onMessage = (msg, contents, actionsDone) => {
         if (Math.random() > probValue.value())
             return false;
 
-        let imageAttachment = msg.attachments.find(v => util.isValidImage(v.filename));
         processFaceSwap(msg, imageAttachment.url).catch(err => console.log(`Failed to run faceapp because ${err}`));
         return true;
     }
@@ -77,6 +141,35 @@ const onMessage = (msg, contents, actionsDone) => {
     return false;
 };
 
+const onDirectMention = (msg, content, actionsDone) => {
+    if (actionsDone)
+        return false;
+
+    let image = msg.attachments.find(v => util.isValidImage(v.filename));
+    if (!image)
+        return false;
+
+    let todayQuota = db.get("faceEditInfo.todayRequests").value();
+    let maxQuota = db.get("faceEditInfo.customRequestsPerDay").value();
+
+    if (todayQuota >= maxQuota) {
+        msg.channel.send(`${msg.author.toString()} Nice image, but I can't do anything to it! (Daily quota hit)`);
+        return true;
+    }
+
+    db.set("faceEditInfo.todayRequests", todayQuota + 1).write();
+
+    processFaceSwap(
+        msg, 
+        image.url,
+        `${msg.author.toString()} Nice image! I don't see anything interesting, though.`,
+        `${msg.author.toString()} ${client.emojis.get("505076258753740810").toString()}`)
+        .catch(err => console.log(`Failed to run faceapp because ${err}`));
+
+    return true;
+};
+
 module.exports = {
-    onMessage: onMessage
+    onMessage: onMessage,
+    onDirectMention: onDirectMention
 };

+ 3 - 0
commands/guide.js

@@ -79,6 +79,9 @@ const onDirectMention = (msg, content, actionsDone) => {
     if (actionsDone)
         return false;
 
+    if(msg.attachments.size > 0)
+        return false;
+
     let parts = content.trim().split(" ");
     let guide = db
         .get("guides")

+ 6 - 0
commands/image_statistics.js

@@ -1,6 +1,7 @@
 const fs = require("fs");
 const util = require("../util.js");
 const path = require("path");
+const db = require("../db.js");
 
 const statsFilePath = path.resolve(path.dirname(module.filename), "../imagestats.csv");
 const statsFile = fs.openSync(statsFilePath, "a");
@@ -11,6 +12,11 @@ const onMessage = msg => {
     if(imagesCount > 0) {
         let now = new Date();
         fs.writeSync(statsFile, `${now.getUTCFullYear()}-${now.getUTCMonth()+1}-${now.getUTCDate()} ${now.getUTCHours()}:${now.getUTCMinutes()};${imagesCount};${msg.channel.name}\n`);
+
+        if(db.get("faceEditStatistics").has(msg.channel.id).value()) {
+            let val = db.get("faceEditStatistics").get(msg.channel.id).value();
+            db.get("faceEditStatistics").set(msg.channel.id, val + 1).write();
+        }
     }
 
     return false;

+ 11 - 0
db.js

@@ -77,6 +77,17 @@ db.defaults({
         "297117726546198528": 0.1,
         "401837400109875216": 0.1
     },
+    faceEditStatistics: {
+        "459622760839118848": 0,
+        "297109482905796608": 0,
+        "297117726546198528": 0,
+        "401837400109875216": 0
+    },
+    faceEditInfo: {
+        quotaRecomputeDate: "",
+        customRequestsPerDay: 0,
+        todayRequests: 0
+    },
     quotes: []
 }).write();
 

+ 1 - 0
package.json

@@ -26,6 +26,7 @@
     "is-image": "^2.0.0",
     "jimp": "^0.5.4",
     "lowdb": "^1.0.0",
+    "node-schedule": "^1.3.0",
     "rss-parser": "^3.4.3",
     "turndown": "^5.0.1",
     "uws": "^99.0.0"