lfs

.gitattributes

@@ -1 +1,2 @@
 hivision/creator/weights/birefnet-v1-lite.onnx filter=lfs diff=lfs merge=lfs -text
+hivision/plugin/font/青鸟华光简琥珀.ttf filter=lfs diff=lfs merge=lfs -text

app.py

@@ -12,6 +12,9 @@ import pathlib
 import numpy as np
 from demo.utils import csv_to_size_list
 import argparse
+from PIL import Image
+from hivision.plugin.watermark import Watermarker, WatermarkerStyles
+
 
 # 获取尺寸列表
 root_dir = os.path.dirname(os.path.abspath(__file__))
@@ -60,6 +63,13 @@ def idphoto_inference(
     custom_image_kb,
     language,
     matting_model_option,
+    watermark_option,
+    watermark_text,
+    watermark_text_color,
+    watermark_text_size,
+    watermark_text_opacity,
+    watermark_text_angle,
+    watermark_text_space,
     face_detect_option,
     head_measure_ratio=0.2,
     # head_height_ratio=0.45,
@@ -234,27 +244,49 @@ def idphoto_inference(
                 )
             )
 
+        # 如果只换底，就不生成排版照
         if (
             idphoto_json["size_mode"]
             == text_lang_map[language]["Only Change Background"]
         ):
             result_layout_image = gr.update(visible=False)
+        # 如果是尺寸列表或自定义尺寸，则生成排版照
         else:
             typography_arr, typography_rotate = generate_layout_photo(
                 input_height=idphoto_json["size"][0],
                 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(
+                            image=result_image_standard,
+                            text=watermark_text,
+                            size=watermark_text_size,
+                            opacity=watermark_text_opacity,
+                            angle=watermark_text_angle,
+                            space=watermark_text_space,
+                            color=watermark_text_color,
+                        ),
+                        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 +304,27 @@ def idphoto_inference(
         else:
             output_image_path = None
 
+        # Add watermark if selected
+        if watermark_option == "添加" or watermark_option == "Add":
+            result_image_standard = add_watermark(
+                image=result_image_standard,
+                text=watermark_text,
+                size=watermark_text_size,
+                opacity=watermark_text_opacity,
+                angle=watermark_text_angle,
+                space=watermark_text_space,
+                color=watermark_text_color,
+            )
+            result_image_hd = add_watermark(
+                image=result_image_hd,
+                text=watermark_text,
+                size=watermark_text_size,
+                opacity=watermark_text_opacity,
+                angle=watermark_text_angle,
+                space=watermark_text_space,
+                color=watermark_text_color,
+            )
+
         if output_image_path:
             result_message = {
                 img_output_standard: result_image_standard,
@@ -292,6 +345,26 @@ def idphoto_inference(
     return result_message
 
 
+# Add this function to handle watermark addition
+def add_watermark(
+    image: np.ndarray, text, size=50, opacity=0.5, angle=45, color="#8B8B1B", space=75
+):
+    image = Image.fromarray(image)
+    # 创建 Watermarker 实例
+    watermarker = Watermarker(
+        input_image=image,
+        text=text,
+        style=WatermarkerStyles.STRIPED,
+        angle=angle,
+        color=color,
+        opacity=opacity,
+        size=size,
+        space=space,
+    )
+    # 返回带水印的图片
+    return np.array(watermarker.image.convert("RGB"))
+
+
 if __name__ == "__main__":
     argparser = argparse.ArgumentParser()
     argparser.add_argument(
@@ -332,6 +405,9 @@ if __name__ == "__main__":
     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 +551,90 @@ if __name__ == "__main__":
                             interactive=True,
                         )
 
+                with gr.Tab("水印") as watermark_parameter_tab:
+                    # 左: 水印
+                    watermark_options = gr.Radio(
+                        choices=watermark_CN,
+                        label="水印",
+                        value="不添加",
+                        elem_id="watermark",
+                    )
+
+                    with gr.Row():
+                        # 左：水印文字
+                        watermark_text_options = gr.Text(
+                            max_length=10,
+                            label="水印文字",
+                            value="Hello",
+                            placeholder="请输入水印文字（最多10个字符）",
+                            elem_id="watermark_text",
+                            interactive=False,
+                        )
+                        # 水印颜色
+                        watermark_text_color = gr.ColorPicker(
+                            label="水印颜色",
+                            interactive=False,
+                            value="#FFFFFF",
+                        )
+
+                    # 文字大小
+                    watermark_text_size = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=20,
+                        label="文字大小",
+                        interactive=False,
+                        step=1,
+                    )
+
+                    # 文字透明度
+                    watermark_text_opacity = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        value=0.15,
+                        label="水印透明度",
+                        interactive=False,
+                        step=0.01,
+                    )
+
+                    # 文字角度
+                    watermark_text_angle = gr.Slider(
+                        minimum=0,
+                        maximum=360,
+                        value=30,
+                        label="水印角度",
+                        interactive=False,
+                        step=1,
+                    )
+
+                    # 文字间距
+                    watermark_text_space = gr.Slider(
+                        minimum=10,
+                        maximum=200,
+                        value=25,
+                        label="水印间距",
+                        interactive=False,
+                        step=1,
+                    )
+
+                    def update_watermark_text_visibility(choice):
+                        return [
+                            gr.update(interactive=(choice == "添加" or choice == "Add"))
+                        ] * 6
+
+                    watermark_options.change(
+                        fn=update_watermark_text_visibility,
+                        inputs=[watermark_options],
+                        outputs=[
+                            watermark_text_options,
+                            watermark_text_color,
+                            watermark_text_size,
+                            watermark_text_opacity,
+                            watermark_text_angle,
+                            watermark_text_space,
+                        ],
+                    )
+
                 img_but = gr.Button("开始制作")
 
                 # 案例图片
@@ -548,6 +708,16 @@ if __name__ == "__main__":
                         top_distance_option: gr.update(label="头距顶距离"),
                         key_parameter_tab: gr.update(label="核心参数"),
                         advance_parameter_tab: gr.update(label="高级参数"),
+                        watermark_parameter_tab: gr.update(label="水印"),
+                        watermark_text_options: gr.update(label="水印文字"),
+                        watermark_text_color: gr.update(label="水印颜色"),
+                        watermark_text_size: gr.update(label="文字大小"),
+                        watermark_text_opacity: gr.update(label="水印透明度"),
+                        watermark_text_angle: gr.update(label="水印角度"),
+                        watermark_text_space: gr.update(label="水印间距"),
+                        watermark_options: gr.update(
+                            label="水印", value="不添加", choices=watermark_CN
+                        ),
                     }
 
                 elif language == "English":
@@ -592,6 +762,16 @@ if __name__ == "__main__":
                         top_distance_option: gr.update(label="Top distance"),
                         key_parameter_tab: gr.update(label="Key Parameters"),
                         advance_parameter_tab: gr.update(label="Advance Parameters"),
+                        watermark_parameter_tab: gr.update(label="Watermark"),
+                        watermark_text_options: gr.update(label="Text"),
+                        watermark_text_color: gr.update(label="Color"),
+                        watermark_text_size: gr.update(label="Size"),
+                        watermark_text_opacity: gr.update(label="Opacity"),
+                        watermark_text_angle: gr.update(label="Angle"),
+                        watermark_text_space: gr.update(label="Space"),
+                        watermark_options: gr.update(
+                            label="Watermark", value="Not Add", choices=watermark_EN
+                        ),
                     }
 
             def change_color(colors):
@@ -652,6 +832,14 @@ if __name__ == "__main__":
                 top_distance_option,
                 key_parameter_tab,
                 advance_parameter_tab,
+                watermark_parameter_tab,
+                watermark_text_options,
+                watermark_text_color,
+                watermark_text_size,
+                watermark_text_opacity,
+                watermark_text_angle,
+                watermark_text_space,
+                watermark_options,
             ],
         )
 
@@ -686,6 +874,13 @@ if __name__ == "__main__":
                 custom_image_kb_size,
                 language_options,
                 matting_model_options,
+                watermark_options,
+                watermark_text_options,
+                watermark_text_color,
+                watermark_text_size,
+                watermark_text_opacity,
+                watermark_text_angle,
+                watermark_text_space,
                 face_detect_model_options,
                 head_measure_ratio_option,
                 top_distance_option,

hivision/.gitattributes

@@ -1,37 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-assets/demoImage.png filter=lfs diff=lfs merge=lfs -text
-hivision/creator/weights/rmbg-1.4.onnx filter=lfs diff=lfs merge=lfs -text

hivision/creator/layout_calculator.py

@@ -128,6 +128,8 @@ def generate_layout_image(
         input_image = cv2.resize(input_image, (width, height))
     if typography_rotate:
         input_image = cv2.transpose(input_image)
+        input_image = cv2.flip(input_image, 0)  # 0 表示垂直镜像
+
         height, width = width, height
     for arr in typography_arr:
         locate_x, locate_y = arr[0], arr[1]

hivision/plugin/font/青鸟华光简琥珀.ttf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d56e625f65cd5c73485e31831a46c612a0b680d2db103886527086fef3d74256
+size 4247512

hivision/plugin/watermark.py

@@ -0,0 +1,232 @@
+"""
+Reference: https://gist.github.com/Deali-Axy/e22ea79bfbe785f9017b2e3cd7fdb3eb
+"""
+
+import enum
+import os
+import math
+import textwrap
+from PIL import Image, ImageFont, ImageDraw, ImageEnhance, ImageChops
+import os
+
+base_path = os.path.abspath(os.path.dirname(__file__))
+
+
+class WatermarkerStyles(enum.Enum):
+    """水印样式"""
+
+    STRIPED = 1  # 斜向重复
+    CENTRAL = 2  # 居中
+
+
+class Watermarker(object):
+    """图片水印工具"""
+
+    def __init__(
+        self,
+        input_image: Image.Image,
+        text: str,
+        style: WatermarkerStyles,
+        angle=30,
+        color="#8B8B1B",
+        font_file="青鸟华光简琥珀.ttf",
+        opacity=0.15,
+        size=50,
+        space=75,
+        chars_per_line=8,
+        font_height_crop=1.2,
+    ):
+        """_summary_
+
+        Parameters
+        ----------
+        input_image : Image.Image
+            PIL图片对象
+        text : str
+            水印文字
+        style : WatermarkerStyles
+            水印样式
+        angle : int, optional
+            水印角度, by default 30
+        color : str, optional
+            水印颜色, by default "#8B8B1B"
+        font_file : str, optional
+            字体文件, by default "青鸟华光简琥珀.ttf"
+        font_height_crop : float, optional
+            字体高度裁剪比例, by default 1.2
+        opacity : float, optional
+            水印透明度, by default 0.15
+        size : int, optional
+            字体大小, by default 50
+        space : int, optional
+            水印间距, by default 75
+        chars_per_line : int, optional
+            每行字符数, by default 8
+        """
+        self.input_image = input_image
+        self.text = text
+        self.style = style
+        self.angle = angle
+        self.color = color
+        self.font_file = os.path.join(base_path, "font", font_file)
+        self.font_height_crop = font_height_crop
+        self.opacity = opacity
+        self.size = size
+        self.space = space
+        self.chars_per_line = chars_per_line
+        self._result_image = None
+
+    @staticmethod
+    def set_image_opacity(image: Image, opacity: float):
+        alpha = image.split()[3]
+        alpha = ImageEnhance.Brightness(alpha).enhance(opacity)
+        image.putalpha(alpha)
+        return image
+
+    @staticmethod
+    def crop_image_edge(image: Image):
+        bg = Image.new(mode="RGBA", size=image.size)
+        diff = ImageChops.difference(image, bg)
+        bbox = diff.getbbox()
+        if bbox:
+            return image.crop(bbox)
+        return image
+
+    def _add_mark_striped(self):
+        origin_image = self.input_image.convert("RGBA")
+        width = len(self.text) * self.size
+        height = round(self.size * self.font_height_crop)
+        watermark_image = Image.new(mode="RGBA", size=(width, height))
+        draw_table = ImageDraw.Draw(watermark_image)
+        draw_table.text(
+            (0, 0),
+            self.text,
+            fill=self.color,
+            font=ImageFont.truetype(self.font_file, size=self.size),
+        )
+        watermark_image = Watermarker.crop_image_edge(watermark_image)
+        Watermarker.set_image_opacity(watermark_image, self.opacity)
+
+        c = int(math.sqrt(origin_image.size[0] ** 2 + origin_image.size[1] ** 2))
+        watermark_mask = Image.new(mode="RGBA", size=(c, c))
+        y, idx = 0, 0
+        while y < c:
+            x = -int((watermark_image.size[0] + self.space) * 0.5 * idx)
+            idx = (idx + 1) % 2
+            while x < c:
+                watermark_mask.paste(watermark_image, (x, y))
+                x += watermark_image.size[0] + self.space
+            y += watermark_image.size[1] + self.space
+
+        watermark_mask = watermark_mask.rotate(self.angle)
+        origin_image.paste(
+            watermark_mask,
+            (int((origin_image.size[0] - c) / 2), int((origin_image.size[1] - c) / 2)),
+            mask=watermark_mask.split()[3],
+        )
+        return origin_image
+
+    def _add_mark_central(self):
+        origin_image = self.input_image.convert("RGBA")
+        text_lines = textwrap.wrap(self.text, width=self.chars_per_line)
+        text = "\n".join(text_lines)
+        width = len(text) * self.size
+        height = round(self.size * self.font_height_crop * len(text_lines))
+        watermark_image = Image.new(mode="RGBA", size=(width, height))
+        draw_table = ImageDraw.Draw(watermark_image)
+        draw_table.text(
+            (0, 0),
+            text,
+            fill=self.color,
+            font=ImageFont.truetype(self.font_file, size=self.size),
+        )
+        watermark_image = Watermarker.crop_image_edge(watermark_image)
+        Watermarker.set_image_opacity(watermark_image, self.opacity)
+
+        c = int(math.sqrt(origin_image.size[0] ** 2 + origin_image.size[1] ** 2))
+        watermark_mask = Image.new(mode="RGBA", size=(c, c))
+        watermark_mask.paste(
+            watermark_image,
+            (
+                int((watermark_mask.width - watermark_image.width) / 2),
+                int((watermark_mask.height - watermark_image.height) / 2),
+            ),
+        )
+        watermark_mask = watermark_mask.rotate(self.angle)
+
+        origin_image.paste(
+            watermark_mask,
+            (
+                int((origin_image.width - watermark_mask.width) / 2),
+                int((origin_image.height - watermark_mask.height) / 2),
+            ),
+            mask=watermark_mask.split()[3],
+        )
+        return origin_image
+
+    @property
+    def image(self):
+        if not self._result_image:
+            if self.style == WatermarkerStyles.STRIPED:
+                self._result_image = self._add_mark_striped()
+            elif self.style == WatermarkerStyles.CENTRAL:
+                self._result_image = self._add_mark_central()
+        return self._result_image
+
+    def save(self, file_path: str, image_format: str = "png"):
+        with open(file_path, "wb") as f:
+            self.image.save(f, image_format)
+
+
+# Gradio 接口
+def watermark_image(
+    image,
+    text,
+    style,
+    angle,
+    color,
+    opacity,
+    size,
+    space,
+):
+    # 创建 Watermarker 实例
+    watermarker = Watermarker(
+        input_image=image,
+        text=text,
+        style=(
+            WatermarkerStyles.STRIPED
+            if style == "STRIPED"
+            else WatermarkerStyles.CENTRAL
+        ),
+        angle=angle,
+        color=color,
+        opacity=opacity,
+        size=size,
+        space=space,
+    )
+
+    # 返回带水印的图片
+    return watermarker.image
+
+
+if __name__ == "__main__":
+    import gradio as gr
+
+    iface = gr.Interface(
+        fn=watermark_image,
+        inputs=[
+            gr.Image(type="pil", label="上传图片", height=400),
+            gr.Textbox(label="水印文字"),
+            gr.Radio(choices=["STRIPED", "CENTRAL"], label="水印样式"),
+            gr.Slider(minimum=0, maximum=360, value=30, label="水印角度"),
+            gr.ColorPicker(label="水印颜色"),
+            gr.Slider(minimum=0, maximum=1, value=0.15, label="水印透明度"),
+            gr.Slider(minimum=10, maximum=100, value=50, label="字体大小"),
+            gr.Slider(minimum=10, maximum=200, value=75, label="水印间距"),
+        ],
+        outputs=gr.Image(type="pil", label="带水印的图片", height=400),
+        title="图片水印工具",
+        description="上传一张图片，添加水印并下载。",
+    )
+
+    iface.launch()