81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
from PIL import Image, ImageDraw, ImageFont
|
|
import numpy as np
|
|
import math
|
|
|
|
# Constants
|
|
DEFAULT_ASPECT_RATIO = 0.6
|
|
FONT_COLOR = (0, 0, 0)
|
|
MARGIN = 1.1
|
|
|
|
|
|
def get_text_square(txt, font_h, aspect_r=DEFAULT_ASPECT_RATIO):
|
|
if not txt:
|
|
return {
|
|
'lineArr': [],
|
|
'numLines': 0,
|
|
'width': 10,
|
|
'height': 10
|
|
}
|
|
|
|
max_line_len = math.ceil(math.sqrt(len(txt)) * 2)
|
|
word_arr = txt.split(' ')
|
|
line_arr = []
|
|
|
|
while word_arr:
|
|
line = ""
|
|
while word_arr and len(line_arr) < max_line_len:
|
|
line += word_arr.pop(0) + ' '
|
|
line_arr.append(line.strip())
|
|
|
|
actual_max_line_len = max(len(s) * 2 for s in line_arr)
|
|
return {
|
|
'lineArr': line_arr,
|
|
'numLines': len(line_arr),
|
|
'width': math.ceil(font_h * aspect_r * actual_max_line_len),
|
|
'height': font_h * len(line_arr)
|
|
}
|
|
|
|
|
|
def write_on_txt_canvas(width, height, txt_to_embed, font_h, margin=MARGIN, font=None):
|
|
img = Image.new('RGBA', (width, height), (255, 255, 255, 0))
|
|
draw = ImageDraw.Draw(img)
|
|
_font_h = int(height * font_h)
|
|
_font = ImageFont.truetype(font, _font_h) if font else None
|
|
text_square = get_text_square(txt_to_embed, _font_h)
|
|
x_offset = 5
|
|
y_offset = 5
|
|
|
|
width_factor = text_square['width'] * margin
|
|
height_factor = text_square['height'] * margin
|
|
|
|
for i in range(int(height / height_factor)):
|
|
for j in range(int(width / width_factor)):
|
|
for k in range(text_square['numLines']):
|
|
x_pos = x_offset + j * width_factor
|
|
y_pos = y_offset + i * height_factor + (k + 1) * _font_h
|
|
draw.text((x_pos, y_pos), text_square['lineArr'][k], fill=FONT_COLOR, font=_font)
|
|
|
|
return np.uint8(img)
|
|
|
|
|
|
def add_watermark(processing_data, txt_data):
|
|
mask = txt_data[:, :, 3] != 0
|
|
|
|
# 显式将数据类型更改为 uint8
|
|
processing_data = processing_data.astype(np.uint8)
|
|
|
|
# 进行位操作
|
|
processing_data[mask, :3] = processing_data[mask, :3] & ~1
|
|
processing_data[~mask, :3] = processing_data[~mask, :3] | 1
|
|
return processing_data
|
|
|
|
|
|
def generate_watermark(input_image, wm_txt, font=None, font_h=0.025):
|
|
text_image = write_on_txt_canvas(input_image.shape[1],
|
|
input_image.shape[0], txt_to_embed=wm_txt,
|
|
font_h=font_h,
|
|
font=font)
|
|
result_img_data = add_watermark(input_image, text_image)
|
|
result_img = Image.fromarray(result_img_data)
|
|
return result_img
|