noctbot/bot/src/plugins/facemorph.ts

import { isValidImage } from "../util";
import Jimp from "jimp";
import { client } from "../client";
import { Message, MessageAttachment } from "discord.js";
import { getRepository } from "typeorm";
import { FaceCaptionMessage, FaceCaptionType } from "@shared/db/entity/FaceCaptionMessage";
import { KnownChannel } from "@shared/db/entity/KnownChannel";
import { logger } from "src/logging";
import got from "got";
import FormData from "form-data";
import { Event, BotEventData, Command, ICommandData, Plugin } from "src/model/plugin";

const EMOTE_GUILD = "505333548694241281";

const CAPTION_IMG_SIZE = 300;
const CAPTION_PROBABILITY = 0.33;

interface Rect {
    x: number,
    y: number,
    w: number,
    h: number
}

interface ErrorInfo {
    ok: boolean,
    error: string;
}

interface FaceData {
    ok: boolean,
    animeFaces: Rect[],
    normalFaces: Rect[]
}

type FaceDetectionResponse = FaceData | ErrorInfo;

function isError(resp: FaceDetectionResponse): resp is ErrorInfo {
    return !resp.ok;
}

type ImageProcessor = (faces: Rect[], data: Buffer) => Promise<Jimp>;
const CAPTION_OFFSET = 5;

@Plugin
export class Facemorph {

    squareFace(rect: Rect): Rect {
        const s = Math.min(rect.w, rect.h);
        return {...rect, w: s, h: s};
    }

    intersects(r1: Rect, r2: Rect): boolean {
        return (
            r1.x <= r2.x + r2.w &&
            r1.x + r1.w >= r2.x &&
            (r1.y <= r2.y + r2.h && r1.y + r1.h >= r2.y)
        );
    }

    morphFaces = async (faces: Rect[], data: Buffer): Promise<Jimp> => {
        const padoru = Math.random() <= this.getPadoruChance();
        let jimpImage = await Jimp.read(data);
        const emoteGuild = client.bot.guilds.resolve(EMOTE_GUILD);
        if (!emoteGuild)
            return jimpImage;
        const emojiKeys = process.env.FOOLS != "TRUE" ? [
            ...emoteGuild
                .emojis.cache.filter(e => !e.animated && e.name?.startsWith("PADORU") == padoru)
                .keys()
        ]: 
            [
                "505335829565276160",
                "430434087157760003",
                "456472341874999297",
                "649677767348060170",
                "589706788782342183",
                "665272109227835422"
            ];

        for (const rect of faces) {
            const dx = rect.x + rect.w / 2;
            const dy = rect.y + rect.h / 2;
            const emojiKey = emojiKeys[Math.floor(Math.random() * emojiKeys.length)];
            const emoji = client.bot.emojis.resolve(emojiKey);
            if (!emoji)
                throw new Error("Failed to resolve emoji!");
            let emojiImage = await Jimp.read(emoji.url);
            let ew = emojiImage.getWidth();
            let eh = emojiImage.getHeight();

            const CONSTANT_SCALE = 1.1;
            const scaleFactor = (Math.max(rect.w, rect.h) / 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;
    }

    async getRandomCaption(type: FaceCaptionType): Promise<FaceCaptionMessage | null> {
        const repo = getRepository(FaceCaptionMessage);
        const caption = await repo.query(`select message
                                        from face_caption_message
                                        where type = $1
                                        order by random()
                                        limit 1`, [type]) as FaceCaptionMessage[];
        if (caption.length == 0)
            return null;
        return caption[0];
    }

    captionFace = async (faces: Rect[], data: Buffer): Promise<Jimp> => {
        const padoru = Math.random() <= this.getPadoruChance();
        const face = faces[Math.floor(Math.random() * faces.length)];
        const squaredFace = this.squareFace(face);
        const targetSize = CAPTION_IMG_SIZE;
        const img = await Jimp.read(data);

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

        const font = await Jimp.loadFont(padoru ? Jimp.FONT_SANS_16_WHITE : Jimp.FONT_SANS_16_BLACK);
        let text = "";
        if(padoru)
            text = "PADORU PADORU";
        else if(process.env.FOOLS == "TRUE") {
            const titles = ["They are horse", "Neigh!", "Insert carrots into them!", "They will become horse!", "They will serve Geoffrey!", "tfw no carrots"];
            text = titles[Math.floor(Math.random() * titles.length)];
        }
        else {
            const prefixMessage = (await this.getRandomCaption(FaceCaptionType.PREFIX))?.message ?? "Feed them";
            const postfixMessage = (await this.getRandomCaption(FaceCaptionType.POSTFIX))?.message ?? "carrots";
            text = `${prefixMessage} ${postfixMessage}`;
        }
        const 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;
    }

    getPadoruChance(): number {
        const now = new Date();
        if (now.getUTCMonth() != 11 || now.getUTCDate() > 25)
            return 0;
        return 1 / (27.0 - now.getUTCDate());
    }

    async processFaceSwap(message: Message, attachmentUrl: string, processor?: ImageProcessor, failMessage?: string, successMessage?: string): Promise<void> {
        const dataResponse = await got.get(attachmentUrl, { responseType: "buffer" });
        if (dataResponse.statusCode != 200) {
            logger.error("Failed to get attachment %s. Got code %s", attachmentUrl, dataResponse.statusCode);
            return;
        }
        const data = dataResponse.body;

        const form = new FormData();
        form.append("img_data", data, {
            filename: "image.png",
            contentType: "application/octet-stream"
        });
        const result = await got.post<FaceDetectionResponse>(`http://${process.env.FACEDETECT_URL}/process`, {
            responseType: "json",
            body: form
        });
        
        if(result.statusCode != 200) {
            logger.error("Face detection failed! Got response %s", result.statusCode);
            return;
        }
        const faceRects = result.body;

        if (isError(faceRects)) {
            logger.error("Face detection failed! Got response %s", result.statusCode);
            return;
        }

        if (faceRects.animeFaces.length == 0 && faceRects.normalFaces.length == 0) {
            if (failMessage) message.channel.send(failMessage);
            return;
        }

        const faces = [...faceRects.normalFaces, ...faceRects.animeFaces];

        let normalCount = faceRects.normalFaces.length;
        let animeCount = faceRects.animeFaces.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 (this.intersects(rAnime, rNormal)) {
                    const animeA = rAnime.w * rAnime.h;
                    const faceA = rNormal.w * rNormal.h;

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

        let jimpImage: Jimp;
        if (processor)
            jimpImage = await processor(faces, data);
        else {
            if (Math.random() <= CAPTION_PROBABILITY)
                jimpImage = await this.captionFace(faces, data);
            else
                jimpImage = await this.morphFaces(faces, data);
        }

        jimpImage.quality(90);
        const buffer = await jimpImage.getBufferAsync(Jimp.MIME_JPEG);
        const messageContents =
            successMessage ||
            `I noticed a face in the image. I think this looks better ${client.bot.emojis.resolve("505076258753740810")?.toString() ?? ":)"}`;

        message.channel.send({
            content: messageContents,
            files: [buffer]
        });
    }

    processLastImage(msg: Message, processor: ImageProcessor): void {
        type AttachedMessage = {msg: Message, att: MessageAttachment};
        const lastImagedMessage = msg.channel.messages.cache.mapValues(m => ({msg: m, att: m.attachments.find(v => isValidImage(v.name))}))
            .filter(v => !v.msg.author.bot && v.att != undefined).last() as AttachedMessage;

        if (!lastImagedMessage) {
            msg.reply("sorry, I couldn't find any recent messages with images.");
            return;
        }

        const replyEmoji = client.bot.emojis.resolve("505076258753740810");
        const emojiText = replyEmoji ? replyEmoji.toString() : "Jiiii~";

        this.processFaceSwap(
            msg,
            lastImagedMessage.att.url,
            processor,
            `${msg.author.toString()} Nice image! I don't see anything interesting, though.`,
            `${msg.author.toString()} ${emojiText}`
        ).catch(err => logger.error(`Failed to run faceapp on message ${msg.id}`, err));
    }

    @Event("message")
    async morphRandomImage(data: BotEventData, msg: Message): Promise<void> {
        if (data.actionsDone) 
            return;

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

        const imageAttachment = msg.attachments.find(v => isValidImage(v.name));

        if (imageAttachment) {
            const repo = getRepository(KnownChannel);

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

            if (!knownChannel || Math.random() > knownChannel.faceMorphProbability)
                return;

            this.processFaceSwap(msg, imageAttachment.url).catch(err =>
                logger.error(`Failed to run faceapp on message ${msg.id}`, err)
            );
            data.actionsDone = true;
        }
    }

    @Event("directMention")
    async morphProvidedImage(data: BotEventData, msg: Message, content: string): Promise<void> {
        if (data.actionsDone) 
            return;

        const image = msg.attachments.find(v => isValidImage(v.name));
        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)`
                );
                data.actionsDone = true;
            }
            return;
        }

        let processor;
        if (content.startsWith("caption this"))
            processor = this.captionFace;
        else
            processor = this.morphFaces;

        const replyEmoji = client.bot.emojis.resolve("505076258753740810");
        const emojiText = replyEmoji ? replyEmoji.toString() : "Jiiii~";

        this.processFaceSwap(
            msg,
            image.url,
            processor,
            `${msg.author.toString()} Nice image! I don't see anything interesting, though.`,
            `${msg.author.toString()} ${emojiText}`
        ).catch(err => logger.error(`Failed to run faceapp because ${msg.id}`, err));

        data.actionsDone = true;
    }

    @Command({
        type: "mention",
        pattern: "caption last image"
    })
    captionLastImage({ message }: ICommandData): void {
        this.processLastImage(message, this.captionFace);
    }

    @Command({
        type: "mention",
        pattern: "look at last image"
    })
    lookLastImage({ message }: ICommandData): void {
        this.processLastImage(message, this.morphFaces);
    }
}