noctbot/commands/facemorph.js

const Clarifai = require("clarifai");
const ClarifaiTOKEN = require("../clarifai_keys.js");
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, failMessage, successMessage) {
    let result = await clarifaiApp.models.predict(Clarifai.FACE_DETECT_MODEL, attachmentUrl);

    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) {
        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();

        const CONSTANT_SCALE = 1.5;
        let scaleFactor = Math.max(bw, bh) / Math.min(ew, eh) * CONSTANT_SCALE;
        ew *= scaleFactor;
        eh *= scaleFactor;

        emojiImage = emojiImage.scale(scaleFactor);
        image = image.composite(emojiImage, dx - ew / 2, dy - eh / 2);
    }

    image.quality(90);
    let buffer = await image.getBufferAsync(Jimp.MIME_JPEG);

    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]
    });
}

const onMessage = (msg, contents, actionsDone) => {
    if (actionsDone)
        return false;

    if(msg.mentions.users.size > 0 && msg.mentions.users.first().id == client.user.id)
        return false;

    let imageAttachment = msg.attachments.find(v => util.isValidImage(v.filename));

    if (imageAttachment) {
        let probValue = db.get("faceEditChannels").get(msg.channel.id);
        if (probValue.isUndefined().value() || probValue.isNull().value())
            return false;

        if (Math.random() > probValue.value())
            return false;

        processFaceSwap(msg, imageAttachment.url).catch(err => console.log(`Failed to run faceapp because ${err}`));
        return true;
    }

    return false;
};

const onDirectMention = (msg, content, actionsDone) => {
    if (actionsDone)
        return false;

    let image = msg.attachments.find(v => util.isValidImage(v.filename));
    if (!image){
        if(msg.attachments.size > 0) {
            msg.channel.send(`${msg.author.toString()} Nice, but I can't do anything to it! (Invalid file type)`);
            return true;
        }
        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,
    onDirectMention: onDirectMention
};