feat: watermark (#83)

* feat:watermark

* ignore

* update ui and fix layout watermark

1. Add watermark as the third tab.
2. Fixed the issue where the layout still had a watermark after choosing not to add a watermark the second time, following the first addition of the watermark.

---------

Co-authored-by: Ze-Yi LIN <[email protected]>

.gitignore

@@ -2,6 +2,7 @@
 **/__pycache__/
 .idea
 .vscode/*
+.history
 .DS_Store
 .env
 demo/kb_output/*.jpg

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,