|
import os |
|
import re |
|
from PIL import Image, ImageDraw, ImageFont |
|
import textwrap |
|
|
|
|
|
def get_unique_filename(filename): |
|
""" |
|
Generate a unique filename by appending a number if a file with the same name already exists. |
|
""" |
|
if not os.path.exists(filename): |
|
return filename |
|
|
|
base, ext = os.path.splitext(filename) |
|
counter = 1 |
|
new_filename = f"{base}_{counter}{ext}" |
|
|
|
while os.path.exists(new_filename): |
|
counter += 1 |
|
new_filename = f"{base}_{counter}{ext}" |
|
|
|
return new_filename |
|
|
|
|
|
def save_image_with_unique_name(image, path): |
|
unique_path = get_unique_filename(path) |
|
image.save(unique_path) |
|
print(f"Image saved as: {unique_path}") |
|
|
|
def find_text_in_answer(text): |
|
print("Full caption:", text) |
|
text = text.split("Caption:")[1] |
|
text = text.replace("\n", "") |
|
text = text.replace("model", "") |
|
|
|
text = re.sub(r'<[^>]*>', '', text) |
|
|
|
|
|
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...") |
|
|
|
|
|
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] |
|
if width <= max_width: |
|
line = test_line |
|
else: |
|
if line: |
|
lines.append(line) |
|
line = word |
|
if line: |
|
lines.append(line) |
|
|
|
y = position[1] |
|
|
|
|
|
for line in lines: |
|
|
|
bbox = draw.textbbox((0, 0), line, font=font) |
|
line_width = bbox[2] - bbox[0] |
|
x = (max_width - line_width) // 2 + position[0] |
|
|
|
|
|
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.text((x, y), line, font=font, fill=text_color) |
|
y += bbox[3] - bbox[1] |
|
|
|
return y - position[1] |
|
|
|
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, output_path, top_margin=10, bottom_margin=10, max_caption_length=10, min_distance_from_bottom_mm=10): |
|
image = Image.open(image_path) |
|
draw = ImageDraw.Draw(image) |
|
width, height = image.size |
|
|
|
|
|
dpi = 96 |
|
min_distance_from_bottom_px = min_distance_from_bottom_mm * dpi / 25.4 |
|
|
|
|
|
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 |
|
|
|
|
|
font = ImageFont.truetype(r"fonts/Anton/Anton-Regular.ttf", font_size) |
|
|
|
|
|
top_caption_position = (width // 10, top_margin) |
|
draw_text(draw, top_caption, top_caption_position, font, width - 2 * (width // 10)) |
|
|
|
|
|
if 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)) |
|
|
|
save_image_with_unique_name(image, output_path) |
|
return image |
|
|
|
|
|
def overlay_caption(text, img_path, output_dir): |
|
img_name = img_path.split("/")[-1] |
|
text = find_text_in_answer(text) |
|
text = text.strip(".") |
|
image = add_caption(img_path, text, output_dir+"/"+img_name) |
|
return image |
|
|