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 morphFaces(faces, data) {
    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);
    }

    return jimpImage;
}

const CAPTION_OFFSET = 5;

async function captionFace(faces, data) {
    let face = faces[Math.floor(Math.random() * faces.length)];
    let squaredFace = await face.toSquareAsync();

    let targetSize = db.get("faceEditConfig.captionedImageSize").value();

    let img = await Jimp.read(data);

    let tempImg = await Jimp.create(squaredFace.width, squaredFace.height);
    tempImg = await tempImg.blit(
        img,
        0,
        0,
        squaredFace.x,
        squaredFace.y,
        squaredFace.width,
        squaredFace.height
    );
    tempImg = await tempImg.scale(targetSize / squaredFace.width);

    let font = await Jimp.loadFont(Jimp.FONT_SANS_16_BLACK);

    let text = `${db.get("faceCaptions.pre").randomElement().value()} ${db.get("faceCaptions.post").randomElement().value()}`;

    let h = Jimp.measureTextHeight(font, text, targetSize - CAPTION_OFFSET * 2);

    let finalImage = await Jimp.create(targetSize, targetSize + h + CAPTION_OFFSET * 2, 0xffffffff);

    finalImage = await finalImage.print(
        font,
        CAPTION_OFFSET,
        CAPTION_OFFSET,
        text,
        finalImage.getWidth() - CAPTION_OFFSET
    );
    finalImage = await finalImage.composite(tempImg, 0, CAPTION_OFFSET * 2 + h);

    return finalImage;
}

async function processFaceSwap(message, attachmentUrl, processor, 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;
    if(processor)
        jimpImage = await processor(faces, data);
    else {
        if(Math.random() <= db.get("faceEditConfig.captionProbability").value())
            jimpImage = await captionFace(faces, data);
        else
            jimpImage = await morphFaces(faces, data);
    }

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

    global.gc();
}

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 processor;
    if(content.startsWith("caption this"))
        processor = captionFace;
    else
        processor = morphFaces;

    processFaceSwap(
        msg,
        image.url,
        processor,
        `${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
};