utils/image_utils.py

import os
import re
from PIL import Image, ImageDraw, ImageFont
import textwrap
import io 
    
def find_text_in_answer(text):
    print("Full caption:", text)
    text = text.split("Caption:")[1]
    text = text.replace("\n", "")
    text = text.replace("model", "")
    # Remove everything that lookslike <>
    text = re.sub(r'<[^>]*>', '', text)
    
    # Remove non-alphanumeric characters (keeping spaces)
    text = re.sub(r'[^a-zA-Z0-9\?\!\s]', '', text)
    print("Filtered caption:", text)
    if text:
        return text
    else:
        return "Me when I couldn't parse the model's answer but I still want you to smile :)"
    
    
def draw_text(draw, text, position, font, max_width, outline_color="black", text_color="white", outline_width=2):
    """
    Draw text on the image with an outline, splitting it into lines if necessary and returning the total height used by the text.
    The text is horizontally centered in the specified max_width.
    """
    print("Adding the caption on the image...")

    # Split the text into multiple lines based on the max width
    lines = []
    words = text.split()
    line = ''
    for word in words:
        test_line = f'{line} {word}'.strip()
        bbox = draw.textbbox((0, 0), test_line, font=font)
        width = bbox[2] - bbox[0]  # Width of the text
        if width <= max_width:
            line = test_line
        else:
            if line:  # Avoid appending empty lines
                lines.append(line)
            line = word
    if line:
        lines.append(line)

    y = position[1]

    # Draw the text with an outline (black) first, centered horizontally
    for line in lines:
        # Calculate the width of the line and adjust the x position to center it
        bbox = draw.textbbox((0, 0), line, font=font)
        line_width = bbox[2] - bbox[0]
        x = (max_width - line_width) // 2 + position[0]

        # Draw the outline by drawing the text multiple times around the original position
        for offset_x in [-outline_width, 0, outline_width]:
            for offset_y in [-outline_width, 0, outline_width]:
                if offset_x != 0 or offset_y != 0:
                    draw.text((x + offset_x, y + offset_y), line, font=font, fill=outline_color)

        # Draw the main text (white) on top of the outline
        draw.text((x, y), line, font=font, fill=text_color)
        y += bbox[3] - bbox[1]  # Update y position based on line height

    return y - position[1]  # Return the total height used by the text

def calculate_text_height(caption, font, max_width):
    """
    Calculate the height of the text when drawn, given the caption, font, and maximum width.
    """
    image = Image.new('RGB', (max_width, 1))
    draw = ImageDraw.Draw(image)
    return draw_text(draw, caption, (0, 0), font, max_width)

def add_caption(image_path, caption, top_margin=10, bottom_margin=10, max_caption_length=10, min_distance_from_bottom_mm=10):
    image = image_path
    draw = ImageDraw.Draw(image)
    width, height = image.size

    # Convert mm to pixels (assuming 96 DPI)
    dpi = 96
    min_distance_from_bottom_px = min_distance_from_bottom_mm * dpi / 25.4

    # Split the caption into two parts if it is too long
    if len(caption.split()) > max_caption_length:
        font_size=20
        total_len = len(caption.split())
        mid = int(total_len / 2)

        top_caption = caption.split()[:mid]
        bottom_caption = caption.split()[mid:]

        top_caption = " ".join(top_caption)
        bottom_caption = " ".join(bottom_caption)
    else:
        top_caption = ""
        bottom_caption = caption
        font_size=30

    # Load a font
    font = ImageFont.truetype(r"fonts/Anton/Anton-Regular.ttf", font_size)

    # Top caption
    top_caption_position = (width // 10, top_margin)
    draw_text(draw, top_caption, top_caption_position, font, width - 2 * (width // 10))

    # Bottom caption
    if bottom_caption:  # Draw bottom caption only if it's not empty
        # Calculate the height of the bottom caption
        bottom_caption_height = calculate_text_height(bottom_caption, font, width - 2 * (width // 10))
        bottom_caption_position = (width // 10, height - min_distance_from_bottom_px - bottom_caption_height)
        draw_text(draw, bottom_caption, bottom_caption_position, font, width - 2 * (width // 10))

    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return buffered.getvalue()
    return image

    
def overlay_caption(text, img_path):
  text = find_text_in_answer(text)
  text = text.strip(".")
  image = add_caption(img_path, text)
  return image