File size: 5,304 Bytes
ceec8fc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
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", "")
# 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, 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
# 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))
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
|