noctbot/commands/facemorph.js

const db = require("../db.js");
const util = require("../util.js");
const Jimp = require("jimp");
const client = require("../client.js");
const cv = require("opencv4nodejs");
const path = require("path");
const request = require("request-promise-native");

const EMOTE_GUILD = "505333548694241281";

const animeCascade = new cv.CascadeClassifier(path.resolve(__dirname, "./animu.xml"));
const faceCascade = new cv.CascadeClassifier(cv.HAAR_FRONTALFACE_ALT2);


function intersects(r1, r2) {
    return (r1.x <= r2.x + r2.width && r1.x + r1.width >= r2.x)
        && (r1.y <= r2.y + r2.height && r1.y + r1.height >= r2.y);
}

async function processFaceSwap(message, attachmentUrl, failMessage, successMessage) {
    let data = await request(attachmentUrl, {encoding: null});

    let im = await cv.imdecodeAsync(data, cv.IMREAD_COLOR);

    let gray = await im.cvtColorAsync(cv.COLOR_BGR2GRAY);
    let normGray = await gray.equalizeHistAsync();

    let animeFaces = await animeCascade.detectMultiScaleAsync(normGray, 1.1, 5, 0, new cv.Size(24, 24));
    let normalFaces = await faceCascade.detectMultiScaleAsync(gray);

    if(animeFaces.objects.length == 0 && normalFaces.objects.length == 0) {
        if (failMessage)
            message.channel.send(failMessage);
        return;
    }
    
    let faces = [...normalFaces.objects, ...animeFaces.objects];

    let normalCount = normalFaces.objects.length;
    let animeCount = animeFaces.objects.length;

    for (let i = 0; i < normalCount; i++) {
        const rNormal = faces[i];
        
        if(animeCount == 0)
            break;

        for (let j = normalCount; j < faces.length; j++) {
            const rAnime = faces[j];
            
            if(intersects(rAnime, rNormal)) {
                let animeA = rAnime.width * rAnime.height;
                let faceA = rNormal.width * rNormal.height;

                if(animeA > faceA){
                    faces.splice(i, 1);
                    normalCount--;
                    i--;
                    break;
                } else {
                    faces.splice(j, 1);
                    animeCount--;
                    j--;
                }
            }
        }
    }

    let jimpImage = await Jimp.read(data);
    let emojiKeys = [...client.guilds.get(EMOTE_GUILD).emojis.filter(e => !e.animated).keys()];

    for (const rect of faces) {
        let dx = rect.x + rect.width / 2;
        let dy = rect.y + rect.height / 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.1;
        let scaleFactor = Math.max(rect.width, rect.height) / Math.min(ew, eh) * CONSTANT_SCALE;
        ew *= scaleFactor;
        eh *= scaleFactor;

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

    jimpImage.quality(90);
    let buffer = await jimpImage.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
};