diff --git a/.gitignore b/.gitignore index e6c38301..7ed25419 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ **/__pycache__/ .idea .vscode/* +.history .DS_Store .env demo/kb_output/*.jpg diff --git a/app.py b/app.py index 7ffc3b5b..06ac1101 100644 --- a/app.py +++ b/app.py @@ -12,6 +12,7 @@ import numpy as np from demo.utils import csv_to_size_list import argparse +from PIL import Image, ImageDraw, ImageFont # 获取尺寸列表 root_dir = os.path.dirname(os.path.abspath(__file__)) @@ -60,6 +61,8 @@ def idphoto_inference( custom_image_kb, language, matting_model_option, + watermark_option, + watermark_text, face_detect_option, head_measure_ratio=0.2, # head_height_ratio=0.45, @@ -245,16 +248,28 @@ def idphoto_inference( input_width=idphoto_json["size"][1], ) - result_layout_image = gr.update( - value=generate_layout_image( - result_image_standard, - typography_arr, - typography_rotate, - height=idphoto_json["size"][0], - width=idphoto_json["size"][1], - ), - visible=True, - ) + if watermark_option == "添加" or watermark_option == "Add": + result_layout_image = gr.update( + value=generate_layout_image( + add_watermark(result_image_standard, watermark_text), + typography_arr, + typography_rotate, + height=idphoto_json["size"][0], + width=idphoto_json["size"][1], + ), + visible=True, + ) + else: + result_layout_image = gr.update( + value=generate_layout_image( + result_image_standard, + typography_arr, + typography_rotate, + height=idphoto_json["size"][0], + width=idphoto_json["size"][1], + ), + visible=True, + ) # 如果输出 KB 大小选择的是自定义 if idphoto_json["custom_image_kb"]: @@ -272,6 +287,11 @@ def idphoto_inference( else: output_image_path = None + # Add watermark if selected + if watermark_option == "添加" or watermark_option == "Add": + result_image_standard = add_watermark(result_image_standard, watermark_text) + result_image_hd = add_watermark(result_image_hd, watermark_text) + if output_image_path: result_message = { img_output_standard: result_image_standard, @@ -292,6 +312,54 @@ def idphoto_inference( return result_message +# Add this function to handle watermark addition +def add_watermark(image, text): + if not text.strip(): + return image + + img = Image.fromarray(image) + img = img.convert("RGBA") + + font_size = 30 + try: + font = ImageFont.truetype("arial.ttf", font_size) + except IOError: + font = ImageFont.load_default() + + # 计算水印文本的大小 + bbox = font.getbbox(text) + text_width = bbox[2] - bbox[0] + text_height = bbox[3] - bbox[1] + + # 创建一个新的图像来绘制倾斜的水印 + watermark = Image.new("RGBA", (img.width * 3, img.height * 3), (0, 0, 0, 0)) + watermark_draw = ImageDraw.Draw(watermark) + + # 在更大的画布上重复绘制水印文本 + for y in range(-img.height, watermark.height, text_height * 3): + for x in range(-img.width, watermark.width, text_width * 3): + watermark_draw.text( + (x, y), text, font=font, fill=(255, 255, 255, 64) + ) # 半透明白色 (255, 255, 255, 64) + + # 旋转水印 + rotated_watermark = watermark.rotate(45, expand=True) + + # 将旋转后的水印裁剪到原始图像大小 + crop_left = (rotated_watermark.width - img.width) // 2 + crop_top = (rotated_watermark.height - img.height) // 2 + crop_right = crop_left + img.width + crop_bottom = crop_top + img.height + cropped_watermark = rotated_watermark.crop( + (crop_left, crop_top, crop_right, crop_bottom) + ) + + # 将水印叠加到原始图像上 + img = Image.alpha_composite(img, cropped_watermark) + + return np.array(img.convert("RGB")) + + if __name__ == "__main__": argparser = argparse.ArgumentParser() argparser.add_argument( @@ -332,6 +400,9 @@ def idphoto_inference( colors_CN = ["蓝色", "白色", "红色", "黑色", "深蓝色", "自定义底色"] colors_EN = ["Blue", "White", "Red", "Black", "Dark blue", "Custom Color"] + watermark_CN = ["不添加", "添加"] + watermark_EN = ["Not Add", "Add"] + render_CN = ["纯色", "上下渐变 (白)", "中心渐变 (白)"] render_EN = ["Solid Color", "Up-Down Gradient (White)", "Center Gradient (White)"] @@ -475,6 +546,33 @@ def load_description(fp): interactive=True, ) + with gr.Tab("水印") as key_parameter_tab: + # 左: 水印 + watermark_options = gr.Radio( + choices=watermark_CN, + label="水印", + value="不添加", + elem_id="watermark", + ) + + # 左:水印文字 + watermark_text_options = gr.Text( + max_length=5, + label="水印文字(最多5)", + value="", + elem_id="watermark_text", + visible=False, + ) + + def update_watermark_text_visibility(choice): + return gr.update(visible=choice == "添加") + + watermark_options.change( + fn=update_watermark_text_visibility, + inputs=[watermark_options], + outputs=[watermark_text_options], + ) + img_but = gr.Button("开始制作") # 案例图片 @@ -686,6 +784,8 @@ def change_image_kb(image_kb_option): custom_image_kb_size, language_options, matting_model_options, + watermark_options, + watermark_text_options, face_detect_model_options, head_measure_ratio_option, top_distance_option,