extensions-builtin/sd-extension-chainner/nodes/impl/caption.py

from __future__ import annotations

import os
import sys
from enum import Enum
from typing import Tuple

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from ..utils.utils import get_h_w_c
from .image_utils import normalize, to_uint8


class CaptionPosition(Enum):
    BOTTOM = "bottom"
    TOP = "top"


def get_font_size(font: ImageFont.FreeTypeFont, text: str) -> Tuple[int, int]:
    """Get font [width, height] of the given text"""
    # (left, top, right, bottom)
    caption_bb = font.getbbox(text)
    font_width = caption_bb[2] - caption_bb[0]
    font_height = caption_bb[3] - caption_bb[1]
    return font_width, font_height


def add_caption(
    img: np.ndarray, caption: str, size: int, position: CaptionPosition
) -> np.ndarray:
    """Add caption with PIL"""
    fontsize = round(size * 0.8)
    if position is CaptionPosition.BOTTOM:
        img = cv2.copyMakeBorder(
            img, 0, size, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0, 1)
        )
    elif position is CaptionPosition.TOP:
        img = cv2.copyMakeBorder(
            img, size, 0, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0, 1)
        )
    else:
        raise RuntimeError(f"Unknown position {position}")

    pimg = Image.fromarray(to_uint8(img))
    font_path = os.path.join(
        os.path.dirname(sys.modules["__main__"].__file__), "fonts/Roboto-Light.ttf"  # type: ignore
    )
    font = ImageFont.truetype(font_path, fontsize)
    h, w, c = get_h_w_c(img)
    text_x = w // 2
    if position is CaptionPosition.BOTTOM:
        text_y = h - round(size / 2)
    elif position is CaptionPosition.TOP:
        text_y = round(size / 2)
    font_color = (255,) * c

    fw, _ = get_font_size(font, caption)
    # scale font size to fit image
    if fw > w:
        font = ImageFont.truetype(font_path, round(fontsize * w / fw))

    d = ImageDraw.Draw(pimg)
    d.text(
        (text_x, text_y),
        caption,
        font=font,
        anchor="mm",
        align="center",
        fill=font_color,
    )

    img = normalize(np.array(pimg))

    return img