noctbot/src/commands/facemorph.ts

import { isValidImage } from "../util";
import Jimp from "jimp";
import { client } from "../client";
import * as cv from "opencv4nodejs";
import * as path from "path";
import request from "request-promise-native";
import { ICommand } from "./command";
import { Message } from "discord.js";
import { getRepository } from "typeorm";
import { FaceCaptionMessage, FaceCaptionType } from "../entity/FaceCaptionMessage";
import { KnownChannel } from "../entity/KnownChannel";

const EMOTE_GUILD = "505333548694241281";

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

const CAPTION_IMG_SIZE = 300;
const CAPTION_PROBABILITY = 0.33;

type ImageProcessor = (faces: cv.Rect[], data: Buffer) => Promise<Jimp>;

function intersects(r1: cv.Rect, r2: cv.Rect) {
    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: cv.Rect[], data: Buffer) {
    let padoru = Math.random() <= getPadoruChance();
    let jimpImage = await Jimp.read(data);
    let emojiKeys = [
        ...client.guilds
            .get(EMOTE_GUILD)
            .emojis.filter(e => !e.animated && e.name.startsWith("PADORU") == padoru)
            .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 getRandomCaption(type: FaceCaptionType) {
    let repo = getRepository(FaceCaptionMessage);

    let caption = await repo.query(`select message
                                    from face_caption_message
                                    where type = ?
                                    order by random()
                                    limit 1`, [ type ]) as FaceCaptionMessage[];
    if(caption.length == 0)
        return null;
    return caption[0];
}

async function captionFace(faces: cv.Rect[], data: Buffer) {
    let padoru = Math.random() <= getPadoruChance();
    let face = faces[Math.floor(Math.random() * faces.length)];
    let squaredFace = await face.toSquareAsync();

    let targetSize = CAPTION_IMG_SIZE;

    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(padoru ? Jimp.FONT_SANS_16_WHITE : Jimp.FONT_SANS_16_BLACK);

    let text = padoru ? "PADORU PADORU" : `${(await getRandomCaption(FaceCaptionType.PREFIX)).message} ${(await getRandomCaption(FaceCaptionType.POSTFIX)).message}`;

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

    let finalImage = await Jimp.create(targetSize, targetSize + h + CAPTION_OFFSET * 2, padoru ? "#FD2027" : "#FFFFFF");

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

/**
 * PADORU PADORU
 */
function getPadoruChance() {
    let now = new Date();
    if (now.getUTCMonth() != 11 || now.getUTCDate() > 25)
        return 0;
    return 1 / (27.0 - now.getUTCDate());
}

async function processFaceSwap(message: Message, attachmentUrl: string, processor?: ImageProcessor, failMessage?: string, successMessage?: string) {
    let data = await request(attachmentUrl, { encoding: null }) as Buffer;

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

function processLastImage(msg: Message, processor: ImageProcessor) {
    let lastImagedMessage = msg.channel.messages.filter(m => !m.author.bot && m.attachments.find(v => isValidImage(v.filename) !== undefined) != undefined).last();

    if (!lastImagedMessage) {
        msg.channel.send(`${msg.author.toString()} Sorry, I couldn't find any recent messages with images.`);
        return;
    }

    let image = lastImagedMessage.attachments.find(v => isValidImage(v.filename));

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

export default {
    onMessage: async (actionsDone, msg, contents) => {
        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 => isValidImage(v.filename));

        if (imageAttachment) {

            let repo = getRepository(KnownChannel);

            let knownChannel = await repo.findOne({
                where: { channelId: msg.channel.id },
                select: [ "faceMorphProbability" ]
            });

            if(!knownChannel || Math.random() > knownChannel.faceMorphProbability)
                return false;
            
            processFaceSwap(msg, imageAttachment.url).catch(err =>
                console.log(`Failed to run faceapp because ${err}`)
            );
            return true;
        }

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

        let image = msg.attachments.find(v => 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;
    },
    commands: [
        {
            pattern: "caption last image",
            action: msg => processLastImage(msg, captionFace)
        },
        {
            pattern: "look at last image",
            action: msg => processLastImage(msg, morphFaces)
        }]
} as ICommand;