goober/modules/image.py

import os
import re
import random
import shutil
import tempfile
from typing import Optional, List
from PIL import Image, ImageDraw, ImageFont, ImageOps
from modules.markovmemory import load_markov_model
from modules.sentenceprocessing import improve_sentence_coherence, rephrase_for_coherence

generated_sentences = set()

def load_font(size):
    return ImageFont.truetype("assets/fonts/Impact.ttf", size=size)
def draw_text_with_outline(draw, text, x, y, font):
    outline_offsets = [(-2, -2), (-2, 2), (2, -2), (2, 2), (0, -2), (0, 2), (-2, 0), (2, 0)]
    for ox, oy in outline_offsets:
        draw.text((x + ox, y + oy), text, font=font, fill="black")
    draw.text((x, y), text, font=font, fill="white")

def fits_in_width(text, font, max_width, draw):
    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    return text_width <= max_width

def split_text_to_fit(text, font, max_width, draw):
    words = text.split()
    for i in range(len(words), 0, -1):
        top_text = " ".join(words[:i])
        bottom_text = " ".join(words[i:])
        if fits_in_width(top_text, font, max_width, draw) and fits_in_width(bottom_text, font, max_width, draw):
            return top_text, bottom_text
    midpoint = len(words) // 2
    return " ".join(words[:midpoint]), " ".join(words[midpoint:])

async def gen_meme(input_image_path, sentence_size=5, max_attempts=10, custom_text=None):
    markov_model = load_markov_model()
    if not markov_model or not os.path.isfile(input_image_path):
        return None

    def generate_text():
        if custom_text:
            return custom_text

        if sentence_size == 1:
            candidate = markov_model.make_short_sentence(max_chars=100, tries=100)
            if candidate:
                candidate = candidate.split()[0]
            return candidate
        else:
            candidate = markov_model.make_sentence(tries=100, max_words=sentence_size)
            if candidate:
                return improve_sentence_coherence(candidate)
            print(candidate)
        return None


    def draw_centered_text(img, text):
        draw = ImageDraw.Draw(img)
        width, height = img.size
        font_size = int(height / 10)
        font = load_font(font_size)

        cleaned = re.sub(r'[^\w\s]', '', text).lower()
        coherent = rephrase_for_coherence(cleaned).upper()

        bbox = draw.textbbox((0, 0), coherent, font=font)
        text_width, text_height_px = bbox[2] - bbox[0], bbox[3] - bbox[1]
        max_text_height = height // 4

        if text_width <= width and text_height_px <= max_text_height:
            draw_text_with_outline(draw, coherent, (width - text_width) / 2, 0, font)
            img.save(input_image_path)
            return True

        top_text, bottom_text = split_text_to_fit(coherent, font, width, draw)
        top_bbox = draw.textbbox((0, 0), top_text, font=font)
        bottom_bbox = draw.textbbox((0, 0), bottom_text, font=font)

        top_height = top_bbox[3] - top_bbox[1]
        bottom_height = bottom_bbox[3] - bottom_bbox[1]

        if top_height <= max_text_height and bottom_height <= max_text_height:
            draw_text_with_outline(draw, top_text, (width - (top_bbox[2] - top_bbox[0])) / 2, 0, font)
            y_bottom = height - bottom_height - int(height * 0.04)
            draw_text_with_outline(draw, bottom_text, (width - (bottom_bbox[2] - bottom_bbox[0])) / 2, y_bottom, font)
            img.save(input_image_path)
            return True

        return False

    attempt = 0
    while attempt < max_attempts:
        response = generate_text() or "NO TEXT GENERATED"
        with Image.open(input_image_path).convert("RGBA") as img:
            if draw_centered_text(img, response):
                return input_image_path
        attempt += 1
    with Image.open(input_image_path).convert("RGBA") as img:
        draw = ImageDraw.Draw(img)
        width, height = img.size
        font_size = int(height / 10)
        font = load_font(font_size)

        truncated = (rephrase_for_coherence(re.sub(r'[^\w\s]', '', "NO TEXT GENERATED").lower()).upper())[:100]
        bbox = draw.textbbox((0, 0), truncated, font=font)
        text_width = bbox[2] - bbox[0]

        draw_text_with_outline(draw, truncated, (width - text_width) / 2, 0, font)
        img.save(input_image_path)
        return input_image_path