Skip to content

Commit f691958

Browse files
AntohhinAntohhin
committed
replace while loop to constant
1 parent 7b42aac commit f691958

File tree

1 file changed

+114
-130
lines changed

1 file changed

+114
-130
lines changed

‎src/storage/watermark.py

Lines changed: 114 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -1,131 +1,115 @@
1-
from pathlib import Path
2-
import random
3-
from io import BytesIO
4-
5-
from PIL import Image, ImageDraw, ImageFont
6-
7-
8-
def draw_text_with_outline(draw, position, text, font, text_colour, outline_colour):
9-
x, y = position
10-
# Draw outline
11-
for adj in range(-1, 2):
12-
for ops in range(-1, 2):
13-
if adj != 0 or ops != 0: # Avoid the center pixel
14-
draw.text((x+adj, y+ops), text, font=font, fill=outline_colour)
15-
draw.text(position, text, font=font, fill=text_colour)
16-
17-
def select_wm_colour(base_brightness) -> tuple:
18-
# if base_brightness > 128:
19-
if base_brightness > 178:
20-
# Black text for lighter background
21-
text_colour = (0, 0, 0, 255)
22-
else:
23-
# White text for darker background
24-
text_colour = (255, 255, 255, 255)
25-
26-
return text_colour
27-
28-
29-
def calculate_corners(img_w, img_h, text_bbox, margin) -> list:
30-
# Estimate text size rely on font and text box
31-
# the (0, 0) is the starting position. return tuple (x1, y1, x2, y2)
32-
33-
text_width = text_bbox[2] - text_bbox[0]
34-
text_height = text_bbox[3] - text_bbox[1]
35-
# Choose a random corner for the text
36-
corners = [
37-
(margin, margin), # Top-left
38-
(img_w - text_width - margin, margin), # Top-right
39-
(margin, img_h - text_height - margin), # Bottom-left
40-
(img_w - text_width - margin, img_h - text_height - margin) # Bottom-right
41-
]
42-
43-
return corners
44-
45-
def check_font_size(text, font_path, image, width_ratio):
46-
# breakpoint = width_ratio * image.size[0]
47-
# fontsize = 20
48-
# learning_rate = 5
49-
fontsize = int(image.size[0] * width_ratio)
50-
print(fontsize)
51-
fontsize = ImageFont.truetype(font_path, fontsize)
52-
53-
# while True:
54-
# if font.getlength(text) < breakpoint:
55-
# fontsize += learning_rate
56-
# else:
57-
# learning_rate = learning_rate // 2
58-
# fontsize -= learning_rate
59-
# font = ImageFont.truetype(font_path, fontsize)
60-
# if learning_rate <= 1:
61-
# break
62-
return fontsize
63-
64-
def draw_corner_watermark(
65-
image_bytes: BytesIO,
66-
text: str,
67-
font_family: str = "Gidole-Regular.ttf",
68-
font_path: str = "static",
69-
margin: int = 24
70-
) -> Image:
71-
72-
with Image.open(image_bytes).convert("RGBA") as base:
73-
txt = Image.new("RGBA", base.size, (255, 255, 255, 0))
74-
75-
d = ImageDraw.Draw(txt)
76-
# Portion of the image the text width should be (between 0 and 1)
77-
width_ratio = .1
78-
# font_path = str(Path(Path.home(), "static", "fonts", font_family))
79-
# font_files_dir = Path(__file__).parent.parent
80-
localization_files_dir = Path(__file__).parent.parent / "static/localization"
81-
print(localization_files_dir, font)
82-
with open(Path(localization_files_dir, "fonts", font_family), "r") as f:
83-
font_path = f
84-
fntsize = check_font_size(text, font_path, base, width_ratio)
85-
font = ImageFont.truetype(font_path, fntsize)
86-
# calculate size of textbox
87-
text_bbox = d.textbbox((0, 0), text, font=font)
88-
# choose a random corner for the text
89-
corners = calculate_corners(img_w=base.size[0], img_h=base.size[1], text_bbox=text_bbox, margin=margin)
90-
text_position = random.choice(corners)
91-
# average brightness of pixel check and switch between black/white
92-
base_brightness = sum(base.getpixel(text_position)[:3]) / 3
93-
text_colour = select_wm_colour(base_brightness)
94-
# define outline colour (opposite of text colour for contrast)
95-
outline_colour = (0, 0, 0, 255) if text_colour == (255, 255, 255, 255) else (255, 255, 255, 255)
96-
draw_text_with_outline(d, text_position, text, fnt, text_colour, outline_colour)
97-
# overlay image of each other
98-
return Image.alpha_composite(base, txt).convert('RGB')
99-
100-
101-
# TODO: async?
102-
def add_watermark(image_content: bytes) -> BytesIO | None:
103-
image_bytes = BytesIO(image_content)
104-
105-
try:
106-
image = draw_corner_watermark(
107-
image_bytes,
108-
text='@ffmemesbot',
109-
# text_size=18,
110-
margin=20
111-
)
112-
except Exception as e:
113-
print(f'Error while adding watermark: {e}')
114-
return None
115-
116-
buff = BytesIO()
117-
buff.name = 'image.jpeg'
118-
image.save(buff, 'JPEG')
119-
buff.seek(0)
120-
121-
return buff
122-
123-
if __name__ == '__main__':
124-
125-
# image_path = Path(Path.home(), "src", "test1.jpeg") # Adjust the path if necessary
126-
image_path = Path(Path.home(), "src", "test2.jpeg") # Adjust the path if necessary
127-
with open(image_path, 'rb') as image_file:
128-
image_bytes = image_file.read()
129-
watermarked_image = add_watermark(image_bytes)
130-
watermarked_image.save('/src/image_Xx.jpg')
1+
from pathlib import Path
2+
import random
3+
from io import BytesIO
4+
5+
from PIL import Image, ImageDraw, ImageFont
6+
7+
8+
def draw_text_with_outline(draw, position, text, font, text_colour, outline_colour):
9+
x, y = position
10+
# Draw outline
11+
for adj in range(-1, 2):
12+
for ops in range(-1, 2):
13+
if adj != 0 or ops != 0: # Avoid the center pixel
14+
draw.text((x+adj, y+ops), text, font=font, fill=outline_colour)
15+
draw.text(position, text, font=font, fill=text_colour)
16+
17+
def select_wm_colour(base_brightness) -> tuple:
18+
# if base_brightness > 128:
19+
if base_brightness > 178:
20+
# Black text for lighter background
21+
text_colour = (0, 0, 0, 255)
22+
else:
23+
# White text for darker background
24+
text_colour = (255, 255, 255, 255)
25+
26+
return text_colour
27+
28+
29+
def calculate_corners(img_w, img_h, text_bbox, margin) -> list:
30+
# Estimate text size rely on font and text box
31+
# the (0, 0) is the starting position. return tuple (x1, y1, x2, y2)
32+
33+
text_width = text_bbox[2] - text_bbox[0]
34+
text_height = text_bbox[3] - text_bbox[1]
35+
# Choose a random corner for the text
36+
corners = [
37+
(margin, margin), # Top-left
38+
(img_w - text_width - margin, margin), # Top-right
39+
(margin, img_h - text_height - margin), # Bottom-left
40+
(img_w - text_width - margin, img_h - text_height - margin) # Bottom-right
41+
]
42+
43+
return corners
44+
45+
def check_font(text, font_path, font_family, image, width_ratio):
46+
# breakpoint = width_ratio * image.size[0]
47+
# fontsize = 20
48+
# learning_rate = 5
49+
fontsize = image.size[0] * width_ratio // 2
50+
font_file = Path(font_path) / font_family # "Gidole-Regular.ttf"
51+
font = ImageFont.truetype(str(font_file), fontsize)
52+
return font
53+
# while True:
54+
# if font.getlength(text) < breakpoint:
55+
# fontsize += learning_rate
56+
# else:
57+
# learning_rate = learning_rate // 2
58+
# fontsize -= learning_rate
59+
# font = ImageFont.truetype(font_path, fontsize)
60+
# if learning_rate <= 1:
61+
# break
62+
# return font
63+
64+
def draw_corner_watermark(
65+
image_bytes: BytesIO,
66+
text: str,
67+
font_family: str = "Gidole-Regular.ttf",
68+
# font_path: str = "static",
69+
margin: int = 20
70+
) -> Image:
71+
72+
with Image.open(image_bytes).convert("RGBA") as base:
73+
txt = Image.new("RGBA", base.size, (255, 255, 255, 0))
74+
75+
d = ImageDraw.Draw(txt)
76+
# ratio of text on the image
77+
width_ratio = .1
78+
fonts_files_dir = Path(__file__).parent.parent / "static/localization/fonts"
79+
font = check_font(text, fonts_files_dir, font_family, base, width_ratio)
80+
# calculate size of textbox
81+
text_bbox = d.textbbox((0, 0), text, font=font)
82+
# choose a random corner for the text
83+
corners = calculate_corners(img_w=base.size[0], img_h=base.size[1], text_bbox=text_bbox, margin=margin)
84+
text_position = random.choice(corners)
85+
# average brightness of pixel check and switch between black/white
86+
base_brightness = sum(base.getpixel(text_position)[:3]) / 3
87+
text_colour = select_wm_colour(base_brightness)
88+
# define outline colour (opposite of text colour for contrast)
89+
outline_colour = (0, 0, 0, 255) if text_colour == (255, 255, 255, 255) else (255, 255, 255, 255)
90+
draw_text_with_outline(d, text_position, text, font, text_colour, outline_colour)
91+
# overlay image of each other
92+
return Image.alpha_composite(base, txt).convert('RGB')
93+
94+
95+
# TODO: async?
96+
def add_watermark(image_content: bytes) -> BytesIO | None:
97+
image_bytes = BytesIO(image_content)
98+
99+
try:
100+
image = draw_corner_watermark(
101+
image_bytes,
102+
text='@ffmemesbot'
103+
)
104+
except Exception as e:
105+
print(f'Error while adding watermark: {e}')
106+
return None
107+
108+
buff = BytesIO()
109+
buff.name = 'image.jpeg'
110+
image.save(buff, 'JPEG')
111+
buff.seek(0)
112+
113+
return buff
114+
131115

0 commit comments

Comments
 (0)