-
Notifications
You must be signed in to change notification settings - Fork 0
/
generator.py
122 lines (103 loc) · 3.61 KB
/
generator.py
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
""" Image generator module """
import os
from typing import List
from PIL import Image, ImageDraw, ImageFont
from PIL.ImageFont import FreeTypeFont
from constants import FOOTER_FONT, FOOTER_TEXT, HEADER_FONT, IMAGE_PADDING
def _wrap_text(font: ImageFont, text: str, max_width: int) -> str:
"""
Add jump lines to wrap the space to the max width.
:param font: ImageFont object.
:param text: Caption str.
:param max_width: Max widht to wrap.
:return: String with the wrapped caption.
"""
if font.getlength(text=text) < max_width:
return text
lines: List[str] = []
words_list: List[str] = text.split(" ")
result: str = ""
result_size: float = 0.0
index: int = 0
while index < len(words_list):
if result_size < max_width:
result += words_list[index] + " "
else:
lines.append(result)
result = words_list[index] + " "
if index == len(words_list) - 1:
lines.append(result)
result_size = font.getlength(text=result)
index += 1
return "\n".join(lines)
def add_caption_to_image(image_name: str, caption: str) -> str:
"""
Generate the image with caption at top.
:param image_name: Name of the target file.
:param caption: Caption to add.
:return: String with the output filename.
"""
# Caption for the image
# Load image
image_path: str = os.path.join(image_name)
image: Image = Image.open(fp=image_path)
caption_font_size: int = int(image.width / 18)
# Create fonts
try:
font: FreeTypeFont = ImageFont.truetype(HEADER_FONT, caption_font_size)
footer_font: FreeTypeFont = ImageFont.truetype(
FOOTER_FONT, int(caption_font_size * 0.5)
)
except Exception:
font: ImageFont = ImageFont.load_default()
footer_font: ImageFont = ImageFont.load_default()
# Split original caption by jump lines
captions: List[str] = []
for sentence in caption.split("\n"):
caption = _wrap_text(
font=font,
text=sentence,
max_width=image.width - caption_font_size * 4.3,
)
captions.append(caption)
caption: str = "\n".join(captions)
# Determine caption lines required
lines: int = len(caption.split("\n")) if "\n" in caption else 1
# Determine blank space to draw caption (margin top)
header_height: int = int(caption_font_size * (lines + 1))
new_height: int = image.height + header_height + int(caption_font_size * 0.6)
# Create new image adding the white space to the target image
new_image: Image = Image.new(
mode=image.mode, size=(image.width + IMAGE_PADDING, new_height), color="white"
)
# Create draw object to insert caption (header) and footer
result: ImageDraw = ImageDraw.Draw(new_image)
result.text(
(caption_font_size, int(caption_font_size * 0.5)),
caption,
font=font,
fill="black",
) # Header
result.text(
(
int(image.width / 2 - footer_font.getlength(text=FOOTER_TEXT) / 2),
image.height + header_height,
),
text=FOOTER_TEXT,
font=footer_font,
fill="black",
) # Footer
# Paste the draw object to the new image
new_image.paste(im=image, box=(int(IMAGE_PADDING / 2), header_height))
# Save result
output_name: str = f"edited-{image_name}"
new_image.save(os.path.join(output_name))
return output_name
def demo():
"""main function"""
add_caption_to_image(
image_name="carbon.png",
caption="Hello world.",
)
if __name__ == "__main__":
demo()