update retinaface

.gitattributes

@@ -3,3 +3,4 @@ hivision/plugin/font/青鸟华光简琥珀.ttf filter=lfs diff=lfs merge=lfs -te
 hivision/creator/weights/hivision_modnet.onnx filter=lfs diff=lfs merge=lfs -text
 hivision/creator/weights/modnet_photographic_portrait_matting.onnx filter=lfs diff=lfs merge=lfs -text
 hivision/creator/weights/rmbg-1.4.onnx filter=lfs diff=lfs merge=lfs -text
+hivision/creator/retinaface/weights/retinaface-resnet50.onnx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -17,6 +17,7 @@ build
 test/temp/*
 !test/temp/.gitkeep
 !hivision/creator/weights/rmbg-1.4.onnx
-!hivision/creator/weights/birefnet-v1-lite.onnx
+!hivision/creator/weights/birefnet-v1-lite.onnx 
+!hivision/creator/retinaface/weights/retinaface-resnet50.onnx
 
 .python-version

app.py

@@ -1,369 +1,32 @@
-import os
-import gradio as gr
-from hivision import IDCreator
-from hivision.error import FaceError, APIError
-from hivision.utils import add_background, resize_image_to_kb
-from hivision.creator.layout_calculator import (
-    generate_layout_photo,
-    generate_layout_image,
-)
-from hivision.creator.choose_handler import choose_handler
-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
-
+import os
+from demo.processor import IDPhotoProcessor
+from demo.ui import create_ui
+from hivision.creator.choose_handler import HUMAN_MATTING_MODELS
 
-# 获取尺寸列表
 root_dir = os.path.dirname(os.path.abspath(__file__))
-size_list_dict_CN = csv_to_size_list(os.path.join(root_dir, "demo/size_list_CN.csv"))
-size_list_dict_EN = csv_to_size_list(os.path.join(root_dir, "demo/size_list_EN.csv"))
-
-color_list_dict_CN = {
-    "蓝色": (86, 140, 212),
-    "白色": (255, 255, 255),
-    "红色": (233, 51, 35),
-    "黑色": (0, 0, 0),
-    "深蓝色": (69, 98, 148),
-}
-
-color_list_dict_EN = {
-    "Blue": (86, 140, 212),
-    "White": (255, 255, 255),
-    "Red": (233, 51, 35),
-    "Black": (0, 0, 0),
-    "Dark blue": (69, 98, 148),
-}
-
-
-# 检测 RGB 是否超出范围，如果超出则约束到 0～255 之间
-def range_check(value, min_value=0, max_value=255):
-    value = int(value)
-    if value <= min_value:
-        value = min_value
-    elif value > max_value:
-        value = max_value
-    return value
-
-
-def idphoto_inference(
-    input_image,
-    mode_option,
-    size_list_option,
-    color_option,
-    render_option,
-    image_kb_options,
-    custom_color_R,
-    custom_color_G,
-    custom_color_B,
-    custom_size_height,
-    custom_size_width,
-    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,
-    top_distance_max=0.12,
-    top_distance_min=0.10,
-):
 
-    top_distance_min = top_distance_max - 0.02
-
-    idphoto_json = {
-        "size_mode": mode_option,
-        "color_mode": color_option,
-        "render_mode": render_option,
-        "image_kb_mode": image_kb_options,
-    }
-
-    text_lang_map = {
-        "中文": {
-            "Size List": "尺寸列表",
-            "Custom Size": "自定义尺寸",
-            "The width should not be greater than the length; the length and width should not be less than 100, and no more than 1800.": "宽度应不大于长度；长宽不应小于 100，大于 1800",
-            "Custom Color": "自定义底色",
-            "Custom": "自定义",
-            "The number of faces is not equal to 1, please upload an image with a single face. If the actual number of faces is 1, it may be an issue with the accuracy of the detection model. Please switch to a different face detection model on the left or raise a Github Issue to notify the author.": "人脸数量不等于 1，请上传单张人脸的图像。如果实际人脸数量为 1，可能是检测模型精度的问题，请在左边更换人脸检测模型或给作者提Github Issue。",
-            "Solid Color": "纯色",
-            "Up-Down Gradient (White)": "上下渐变 (白)",
-            "Center Gradient (White)": "中心渐变 (白)",
-            "Set KB size (Download in the bottom right)": "设置 KB 大小（结果在右边最底的组件下载）",
-            "Not Set": "不设置",
-            "Only Change Background": "只换底",
-        },
-        "English": {
-            "Size List": "Size List",
-            "Custom Size": "Custom Size",
-            "The width should not be greater than the length; the length and width should not be less than 100, and no more than 1800.": "The width should not be greater than the length; the length and width should not be less than 100, and no more than 1800.",
-            "Custom Color": "Custom Color",
-            "Custom": "Custom",
-            "The number of faces is not equal to 1, please upload an image with a single face. If the actual number of faces is 1, it may be an issue with the accuracy of the detection model. Please switch to a different face detection model on the left or raise a Github Issue to notify the author.": "The number of faces is not equal to 1, please upload an image with a single face. If the actual number of faces is 1, it may be an issue with the accuracy of the detection model. Please switch to a different face detection model on the left or raise a Github Issue to notify the author.",
-            "Solid Color": "Solid Color",
-            "Up-Down Gradient (White)": "Up-Down Gradient (White)",
-            "Center Gradient (White)": "Center Gradient (White)",
-            "Set KB size (Download in the bottom right)": "Set KB size (Download in the bottom right)",
-            "Not Set": "Not Set",
-            "Only Change Background": "Only Change Background",
-        },
-    }
-
-    # 如果尺寸模式选择的是尺寸列表
-    if idphoto_json["size_mode"] == text_lang_map[language]["Size List"]:
-        if language == "中文":
-            idphoto_json["size"] = size_list_dict_CN[size_list_option]
-        else:
-            idphoto_json["size"] = size_list_dict_EN[size_list_option]
-    # 如果尺寸模式选择的是自定义尺寸
-    elif idphoto_json["size_mode"] == text_lang_map[language]["Custom Size"]:
-        id_height = int(custom_size_height)
-        id_width = int(custom_size_width)
-        if (
-            id_height < id_width
-            or min(id_height, id_width) < 100
-            or max(id_height, id_width) > 1800
-        ):
-            return {
-                img_output_standard: gr.update(value=None),
-                img_output_standard_hd: gr.update(value=None),
-                notification: gr.update(
-                    value=text_lang_map[language][
-                        "The width should not be greater than the length; the length and width should not be less than 100, and no more than 1800."
-                    ],
-                    visible=True,
-                ),
-            }
-        idphoto_json["size"] = (id_height, id_width)
-    else:
-        idphoto_json["size"] = (None, None)
-
-    # 如果颜色模式选择的是自定义底色
-    if idphoto_json["color_mode"] == text_lang_map[language]["Custom Color"]:
-        idphoto_json["color_bgr"] = (
-            range_check(custom_color_R),
-            range_check(custom_color_G),
-            range_check(custom_color_B),
-        )
-    else:
-        if language == "中文":
-            idphoto_json["color_bgr"] = color_list_dict_CN[color_option]
-        else:
-            idphoto_json["color_bgr"] = color_list_dict_EN[color_option]
-
-    # 如果输出 KB 大小选择的是自定义
-    if idphoto_json["image_kb_mode"] == text_lang_map[language]["Custom"]:
-        idphoto_json["custom_image_kb"] = custom_image_kb
-    else:
-        idphoto_json["custom_image_kb"] = None
-
-    creator = IDCreator()
-    choose_handler(creator, matting_model_option, face_detect_option)
-
-    change_bg_only = idphoto_json["size_mode"] in ["只换底", "Only Change Background"]
-    # 生成证件照
-    try:
-        result = creator(
-            input_image,
-            change_bg_only=change_bg_only,
-            size=idphoto_json["size"],
-            head_measure_ratio=head_measure_ratio,
-            # head_height_ratio=head_height_ratio,
-            head_top_range=(top_distance_max, top_distance_min),
+HUMAN_MATTING_MODELS_EXIST = [
+    os.path.splitext(file)[0]
+    for file in os.listdir(os.path.join(root_dir, "hivision/creator/weights"))
+    if file.endswith(".onnx") or file.endswith(".mnn")
+]
+# 在HUMAN_MATTING_MODELS中的模型才会被加载到Gradio中显示
+HUMAN_MATTING_MODELS = [
+    model for model in HUMAN_MATTING_MODELS if model in HUMAN_MATTING_MODELS_EXIST
+]
+
+FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn"]
+FACE_DETECT_MODELS_EXPAND = (
+    ["retinaface-resnet50"]
+    if os.path.exists(
+        os.path.join(
+            root_dir, "hivision/creator/retinaface/weights/retinaface-resnet50.onnx"
         )
-    except FaceError:
-        result_message = {
-            img_output_standard: gr.update(value=None),
-            img_output_standard_hd: gr.update(value=None),
-            img_output_layout: gr.update(visible=False),
-            notification: gr.update(
-                value=text_lang_map[language][
-                    "The number of faces is not equal to 1, please upload an image with a single face. If the actual number of faces is 1, it may be an issue with the accuracy of the detection model. Please switch to a different face detection model on the left or raise a Github Issue to notify the author."
-                ],
-                visible=True,
-            ),
-        }
-    except APIError as e:
-        result_message = {
-            img_output_standard: gr.update(value=None),
-            img_output_standard_hd: gr.update(value=None),
-            img_output_layout: gr.update(visible=False),
-            notification: gr.update(
-                value=f"Please make sure you have correctly set up the Face++ API Key and Secret.\nAPI Error\nStatus Code is {e.status_code}\nPossible errors are: {e}\n",
-                visible=True,
-            ),
-        }
-    else:
-        (result_image_standard, result_image_hd, _, _) = result
-        if idphoto_json["render_mode"] == text_lang_map[language]["Solid Color"]:
-            result_image_standard = np.uint8(
-                add_background(result_image_standard, bgr=idphoto_json["color_bgr"])
-            )
-            result_image_hd = np.uint8(
-                add_background(result_image_hd, bgr=idphoto_json["color_bgr"])
-            )
-        elif (
-            idphoto_json["render_mode"]
-            == text_lang_map[language]["Up-Down Gradient (White)"]
-        ):
-            result_image_standard = np.uint8(
-                add_background(
-                    result_image_standard,
-                    bgr=idphoto_json["color_bgr"],
-                    mode="updown_gradient",
-                )
-            )
-            result_image_hd = np.uint8(
-                add_background(
-                    result_image_hd,
-                    bgr=idphoto_json["color_bgr"],
-                    mode="updown_gradient",
-                )
-            )
-        else:
-            result_image_standard = np.uint8(
-                add_background(
-                    result_image_standard,
-                    bgr=idphoto_json["color_bgr"],
-                    mode="center_gradient",
-                )
-            )
-            result_image_hd = np.uint8(
-                add_background(
-                    result_image_hd,
-                    bgr=idphoto_json["color_bgr"],
-                    mode="center_gradient",
-                )
-            )
-
-        # 如果只换底，就不生成排版照
-        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],
-            )
-
-            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"]:
-            # 将标准照大小调整至目标大小
-            print("调整 kb 大小到", idphoto_json["custom_image_kb"], "kb")
-            # 输出路径为一个根据时间戳 + 哈希值生成的随机文件名
-            import time
-
-            output_image_path = f"{os.path.join(os.path.dirname(__file__), 'demo/kb_output')}/{int(time.time())}.jpg"
-            resize_image_to_kb(
-                result_image_standard,
-                output_image_path,
-                idphoto_json["custom_image_kb"],
-            )
-        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,
-                img_output_standard_hd: result_image_hd,
-                img_output_layout: result_layout_image,
-                notification: gr.update(visible=False),
-                file_download: gr.update(visible=True, value=output_image_path),
-            }
-        else:
-            result_message = {
-                img_output_standard: result_image_standard,
-                img_output_standard_hd: result_image_hd,
-                img_output_layout: result_layout_image,
-                notification: gr.update(visible=False),
-                file_download: gr.update(visible=False),
-            }
-
-    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"))
-
+    else []
+)
+FACE_DETECT_MODELS += FACE_DETECT_MODELS_EXPAND
 
 if __name__ == "__main__":
     argparser = argparse.ArgumentParser()
@@ -379,525 +42,17 @@ if __name__ == "__main__":
         default=None,
         help="The root path of the server, default is None (='/'), e.g. '/myapp'",
     )
-
     args = argparser.parse_args()
 
-    language = ["中文", "English"]
-
-    matting_model_list = [
-        os.path.splitext(file)[0]
-        for file in os.listdir(os.path.join(root_dir, "hivision/creator/weights"))
-        if file.endswith(".onnx") or file.endswith(".mnn")
-    ]
-    DEFAULT_MATTING_MODEL = "modnet_photographic_portrait_matting"
-    if DEFAULT_MATTING_MODEL in matting_model_list:
-        matting_model_list.remove(DEFAULT_MATTING_MODEL)
-        matting_model_list.insert(0, DEFAULT_MATTING_MODEL)
-
-    face_detect_model_list = ["mtcnn", "face++ (联网API)"]
-
-    size_mode_CN = ["尺寸列表", "只换底", "自定义尺寸"]
-    size_mode_EN = ["Size List", "Only Change Background", "Custom Size"]
-
-    size_list_CN = list(size_list_dict_CN.keys())
-    size_list_EN = list(size_list_dict_EN.keys())
-
-    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)"]
-
-    image_kb_CN = ["不设置", "自定义"]
-    image_kb_EN = ["Not Set", "Custom"]
-
-    css = """
-        #col-left {
-            margin: 0 auto;
-            max-width: 430px;
-        }
-        #col-mid {
-            margin: 0 auto;
-            max-width: 430px;
-        }
-        #col-right {
-            margin: 0 auto;
-            max-width: 430px;
-        }
-        #col-showcase {
-            margin: 0 auto;
-            max-width: 1100px;
-        }
-        #button {
-            color: blue;
-        }
-        """
-
-    def load_description(fp):
-        with open(fp, "r", encoding="utf-8") as f:
-            content = f.read()
-        return content
-
-    demo = gr.Blocks(title="HivisionIDPhotos", css=css)
-
-    with demo:
-        gr.HTML(load_description(os.path.join(root_dir, "assets/title.md")))
-        with gr.Row():
-            # ------------ 左半边 UI ----------------
-            with gr.Column():
-                img_input = gr.Image(height=400)
-
-                with gr.Row():
-                    language_options = gr.Dropdown(
-                        choices=language,
-                        label="Language",
-                        value="中文",
-                        elem_id="language",
-                    )
-                    face_detect_model_options = gr.Dropdown(
-                        choices=face_detect_model_list,
-                        label="人脸检测模型",
-                        value=face_detect_model_list[0],
-                        elem_id="matting_model",
-                    )
-                    matting_model_options = gr.Dropdown(
-                        choices=matting_model_list,
-                        label="抠图模型",
-                        value=matting_model_list[0],
-                        elem_id="matting_model",
-                    )
-
-                with gr.Tab("核心参数") as key_parameter_tab:
-                    mode_options = gr.Radio(
-                        choices=size_mode_CN,
-                        label="证件照尺寸选项",
-                        value="尺寸列表",
-                        elem_id="size",
-                    )
-
-                    # 预设尺寸下拉菜单
-                    with gr.Row(visible=True) as size_list_row:
-                        size_list_options = gr.Dropdown(
-                            choices=size_list_CN,
-                            label="预设尺寸",
-                            value=size_list_CN[0],
-                            elem_id="size_list",
-                        )
-
-                    with gr.Row(visible=False) as custom_size:
-                        custom_size_height = gr.Number(
-                            value=413, label="height", interactive=True
-                        )
-                        custom_size_wdith = gr.Number(
-                            value=295, label="width", interactive=True
-                        )
-
-                    # 左：背景色选项
-                    color_options = gr.Radio(
-                        choices=colors_CN, label="背景色", value="蓝色", elem_id="color"
-                    )
-
-                    # 左：如果选择「自定义底色」，显示 RGB 输入框
-                    with gr.Row(visible=False) as custom_color:
-                        custom_color_R = gr.Number(value=0, label="R", interactive=True)
-                        custom_color_G = gr.Number(value=0, label="G", interactive=True)
-                        custom_color_B = gr.Number(value=0, label="B", interactive=True)
-
-                    # 左：渲染方式选项
-                    render_options = gr.Radio(
-                        choices=render_CN,
-                        label="渲染方式",
-                        value="纯色",
-                        elem_id="render",
-                    )
-
-                with gr.Tab("高级参数") as advance_parameter_tab:
-                    # 面部占照片总比例
-                    head_measure_ratio_option = gr.Slider(
-                        minimum=0.1,
-                        maximum=0.5,
-                        value=0.2,
-                        step=0.01,
-                        label="面部比例",
-                        interactive=True,
-                    )
-                    # 人像头顶距离照片顶部的比例
-                    top_distance_option = gr.Slider(
-                        minimum=0.02,
-                        maximum=0.5,
-                        value=0.12,
-                        step=0.01,
-                        label="头距顶距离",
-                        interactive=True,
-                    )
-
-                    # 输出照片的KB值
-                    image_kb_options = gr.Radio(
-                        choices=image_kb_CN,
-                        label="设置 KB 大小（结果在右边最底的组件下载）",
-                        value="不设置",
-                        elem_id="image_kb",
-                    )
-                    # 自定义 KB 大小，滑动条，最小 10KB，最大 200KB
-                    with gr.Row(visible=False) as custom_image_kb:
-                        custom_image_kb_size = gr.Slider(
-                            minimum=10,
-                            maximum=1000,
-                            value=50,
-                            label="KB 大小",
-                            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("开始制作")
-
-                # 案例图片
-                example_images = gr.Examples(
-                    inputs=[img_input],
-                    examples=[
-                        [path.as_posix()]
-                        for path in sorted(
-                            pathlib.Path(os.path.join(root_dir, "demo/images")).rglob(
-                                "*.jpg"
-                            )
-                        )
-                    ],
-                )
-
-            # ---------------- 右半边 UI ----------------
-            with gr.Column():
-                notification = gr.Text(label="状态", visible=False)
-                with gr.Row():
-                    img_output_standard = gr.Image(
-                        label="标准照", height=350, format="jpeg"
-                    )
-                    img_output_standard_hd = gr.Image(
-                        label="高清照", height=350, format="jpeg"
-                    )
-                img_output_layout = gr.Image(
-                    label="六寸排版照", height=350, format="jpeg"
-                )
-                file_download = gr.File(label="下载调整 KB 大小后的照片", visible=False)
-
-            # ---------------- 设置隐藏/显示组件 ----------------
-            def change_language(language):
-                # 将Gradio组件中的内容改为中文或英文
-                if language == "中文":
-                    return {
-                        size_list_options: gr.update(
-                            label="预设尺寸",
-                            choices=size_list_CN,
-                            value=size_list_CN[0],
-                        ),
-                        mode_options: gr.update(
-                            label="证件照尺寸选项",
-                            choices=size_mode_CN,
-                            value="尺寸列表",
-                        ),
-                        color_options: gr.update(
-                            label="背景色",
-                            choices=colors_CN,
-                            value="蓝色",
-                        ),
-                        img_but: gr.update(value="开始制作"),
-                        render_options: gr.update(
-                            label="渲染方式",
-                            choices=render_CN,
-                            value="纯色",
-                        ),
-                        image_kb_options: gr.update(
-                            label="设置 KB 大小（结果在右边最底的组件下载）",
-                            choices=image_kb_CN,
-                            value="不设置",
-                        ),
-                        matting_model_options: gr.update(label="抠图模型"),
-                        face_detect_model_options: gr.update(label="人脸检测模型"),
-                        custom_image_kb_size: gr.update(label="KB 大小"),
-                        notification: gr.update(label="状态"),
-                        img_output_standard: gr.update(label="标准照"),
-                        img_output_standard_hd: gr.update(label="高清照"),
-                        img_output_layout: gr.update(label="六寸排版照"),
-                        file_download: gr.update(label="下载调整 KB 大小后的照片"),
-                        head_measure_ratio_option: gr.update(label="面部比例"),
-                        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":
-                    return {
-                        size_list_options: gr.update(
-                            label="Default size",
-                            choices=size_list_EN,
-                            value=size_list_EN[0],
-                        ),
-                        mode_options: gr.update(
-                            label="ID photo size options",
-                            choices=size_mode_EN,
-                            value="Size List",
-                        ),
-                        color_options: gr.update(
-                            label="Background color",
-                            choices=colors_EN,
-                            value="Blue",
-                        ),
-                        img_but: gr.update(value="Start"),
-                        render_options: gr.update(
-                            label="Rendering mode",
-                            choices=render_EN,
-                            value="Solid Color",
-                        ),
-                        image_kb_options: gr.update(
-                            label="Set KB size (Download in the bottom right)",
-                            choices=image_kb_EN,
-                            value="Not Set",
-                        ),
-                        matting_model_options: gr.update(label="Matting model"),
-                        face_detect_model_options: gr.update(label="Face detect model"),
-                        custom_image_kb_size: gr.update(label="KB size"),
-                        notification: gr.update(label="Status"),
-                        img_output_standard: gr.update(label="Standard photo"),
-                        img_output_standard_hd: gr.update(label="HD photo"),
-                        img_output_layout: gr.update(label="Layout photo"),
-                        file_download: gr.update(
-                            label="Download the photo after adjusting the KB size"
-                        ),
-                        head_measure_ratio_option: gr.update(label="Head ratio"),
-                        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):
-                if colors == "自定义底色" or colors == "Custom Color":
-                    return {custom_color: gr.update(visible=True)}
-                else:
-                    return {custom_color: gr.update(visible=False)}
-
-            def change_size_mode(size_option_item):
-                if (
-                    size_option_item == "自定义尺寸"
-                    or size_option_item == "Custom Size"
-                ):
-                    return {
-                        custom_size: gr.update(visible=True),
-                        size_list_row: gr.update(visible=False),
-                    }
-                elif (
-                    size_option_item == "只换底"
-                    or size_option_item == "Only Change Background"
-                ):
-                    return {
-                        custom_size: gr.update(visible=False),
-                        size_list_row: gr.update(visible=False),
-                    }
-                else:
-                    return {
-                        custom_size: gr.update(visible=False),
-                        size_list_row: gr.update(visible=True),
-                    }
-
-            def change_image_kb(image_kb_option):
-                if image_kb_option == "自定义" or image_kb_option == "Custom":
-                    return {custom_image_kb: gr.update(visible=True)}
-                else:
-                    return {custom_image_kb: gr.update(visible=False)}
-
-        # ---------------- 绑定事件 ----------------
-        language_options.input(
-            change_language,
-            inputs=[language_options],
-            outputs=[
-                size_list_options,
-                mode_options,
-                color_options,
-                img_but,
-                render_options,
-                image_kb_options,
-                matting_model_options,
-                face_detect_model_options,
-                custom_image_kb_size,
-                notification,
-                img_output_standard,
-                img_output_standard_hd,
-                img_output_layout,
-                file_download,
-                head_measure_ratio_option,
-                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,
-            ],
-        )
-
-        color_options.input(
-            change_color, inputs=[color_options], outputs=[custom_color]
-        )
-
-        mode_options.input(
-            change_size_mode,
-            inputs=[mode_options],
-            outputs=[custom_size, size_list_row],
-        )
-
-        image_kb_options.input(
-            change_image_kb, inputs=[image_kb_options], outputs=[custom_image_kb]
-        )
-
-        img_but.click(
-            idphoto_inference,
-            inputs=[
-                img_input,
-                mode_options,
-                size_list_options,
-                color_options,
-                render_options,
-                image_kb_options,
-                custom_color_R,
-                custom_color_G,
-                custom_color_B,
-                custom_size_height,
-                custom_size_wdith,
-                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,
-            ],
-            outputs=[
-                img_output_standard,
-                img_output_standard_hd,
-                img_output_layout,
-                notification,
-                file_download,
-            ],
-        )
+    processor = IDPhotoProcessor()
 
+    demo = create_ui(
+        processor, root_dir, HUMAN_MATTING_MODELS_EXIST, FACE_DETECT_MODELS
+    )
     demo.launch(
         # server_name=args.host,
         # server_port=args.port,
-        show_api=False,
         # favicon_path=os.path.join(root_dir, "assets/hivision_logo.png"),
         # root_path=args.root_path,
+        show_api=False,
     )

demo/assets/color_list_CN.csv

@@ -0,0 +1,6 @@
+Name,Hex
+蓝色,628bce
+白色,ffffff
+红色,d74532
+黑色,000000
+深蓝色,4b6190

demo/assets/color_list_EN.csv

@@ -0,0 +1,6 @@
+Name,Hex  
+Blue,628bce  
+White,ffffff  
+Red,d74532  
+Black,000000  
+Dark Blue,4b6190

demo/config.py

@@ -0,0 +1,19 @@
+import os
+from demo.utils import csv_to_size_list, csv_to_color_list
+
+
+def load_configuration(root_dir):
+    size_list_dict_CN = csv_to_size_list(
+        os.path.join(root_dir, "assets/size_list_CN.csv")
+    )
+    size_list_dict_EN = csv_to_size_list(
+        os.path.join(root_dir, "assets/size_list_EN.csv")
+    )
+    color_list_dict_CN = csv_to_color_list(
+        os.path.join(root_dir, "assets/color_list_CN.csv")
+    )
+    color_list_dict_EN = csv_to_color_list(
+        os.path.join(root_dir, "assets/color_list_EN.csv")
+    )
+
+    return size_list_dict_CN, size_list_dict_EN, color_list_dict_CN, color_list_dict_EN

demo/locals.py

@@ -0,0 +1,271 @@
+# Copyright 2024 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from demo.utils import csv_to_size_list
+from demo.config import load_configuration
+import os
+
+base_dir = os.path.dirname(os.path.abspath(__file__))
+size_list_dict_CN = csv_to_size_list(os.path.join(base_dir, "assets/size_list_CN.csv"))
+size_list_dict_EN = csv_to_size_list(os.path.join(base_dir, "assets/size_list_EN.csv"))
+(
+    size_list_config_CN,
+    size_list_config_EN,
+    color_list_dict_CN,
+    color_list_dict_EN,
+) = load_configuration(base_dir)
+
+
+LOCALES = {
+    "face_model": {
+        "en": {
+            "label": "Face detection model",
+        },
+        "zh": {
+            "label": "人脸检测模型",
+        },
+    },
+    "matting_model": {
+        "en": {
+            "label": "Matting model",
+        },
+        "zh": {
+            "label": "抠图模型",
+        },
+    },
+    "key_param": {
+        "en": {
+            "label": "Key Parameters",
+        },
+        "zh": {
+            "label": "核心参数",
+        },
+    },
+    "advance_param": {
+        "en": {
+            "label": "Advance Parameters",
+        },
+        "zh": {
+            "label": "高级参数",
+        },
+    },
+    "size_mode": {
+        "en": {
+            "label": "ID photo size options",
+            "choices": ["Size List", "Only Change Background", "Custom Size"],
+            "custom_size_eror": "The width should not be greater than the length; the length and width should not be less than 100, and no more than 1800.",
+        },
+        "zh": {
+            "label": "证件照尺寸选项",
+            "choices": ["尺寸列表", "只换底", "自定义尺寸"],
+            "custom_size_eror": "宽度不应大于长度；长度和宽度不应小于100，不大于1800。",
+        },
+    },
+    "size_list": {
+        "en": {
+            "label": "Size list",
+            "choices": list(size_list_dict_EN.keys()),
+            "develop": size_list_config_EN,
+        },
+        "zh": {
+            "label": "预设尺寸",
+            "choices": list(size_list_dict_CN.keys()),
+            "develop": size_list_config_CN,
+        },
+    },
+    "bg_color": {
+        "en": {
+            "label": "Background color",
+            "choices": list(color_list_dict_EN.keys()) + ["Custom"],
+            "develop": color_list_dict_EN,
+        },
+        "zh": {
+            "label": "背景颜色",
+            "choices": list(color_list_dict_CN.keys()) + ["自定义底色"],
+            "develop": color_list_dict_CN,
+        },
+    },
+    "button": {
+        "en": {
+            "label": "Start",
+        },
+        "zh": {
+            "label": "开始制作",
+        },
+    },
+    "head_measure_ratio": {
+        "en": {
+            "label": "Head ratio",
+        },
+        "zh": {
+            "label": "面部比例",
+        },
+    },
+    "top_distance": {
+        "en": {
+            "label": "Top distance",
+        },
+        "zh": {
+            "label": "头距顶距离",
+        },
+    },
+    "image_kb": {
+        "en": {
+            "label": "Set KB size",
+            "choices": ["Not Set", "Custom"],
+        },
+        "zh": {
+            "label": "设置 KB 大小",
+            "choices": ["不设置", "自定义"],
+        },
+    },
+    "image_kb_size": {
+        "en": {
+            "label": "KB size",
+        },
+        "zh": {
+            "label": "KB 大小",
+        },
+    },
+    "render_mode": {
+        "en": {
+            "label": "Render mode",
+            "choices": [
+                "Solid Color",
+                "Up-Down Gradient (White)",
+                "Center Gradient (White)",
+            ],
+        },
+        "zh": {
+            "label": "渲染方式",
+            "choices": ["纯色", "上下渐变（白色）", "中心渐变（白色）"],
+        },
+    },
+    # Tab3 - 水印工作台
+    "watermark_tab": {
+        "en": {
+            "label": "Watermark",
+        },
+        "zh": {
+            "label": "水印",
+        },
+    },
+    "watermark_text": {
+        "en": {
+            "label": "Text",
+            "value": "Hello",
+            "placeholder": "up to 20 characters",
+        },
+        "zh": {
+            "label": "水印文字",
+            "value": "Hello",
+            "placeholder": "最多20个字符",
+        },
+    },
+    "watermark_color": {
+        "en": {
+            "label": "Color",
+        },
+        "zh": {
+            "label": "水���颜色",
+        },
+    },
+    "watermark_size": {
+        "en": {
+            "label": "Size",
+        },
+        "zh": {
+            "label": "文字大小",
+        },
+    },
+    "watermark_opacity": {
+        "en": {
+            "label": "Opacity",
+        },
+        "zh": {
+            "label": "水印透明度",
+        },
+    },
+    "watermark_angle": {
+        "en": {
+            "label": "Angle",
+        },
+        "zh": {
+            "label": "水印角度",
+        },
+    },
+    "watermark_space": {
+        "en": {
+            "label": "Space",
+        },
+        "zh": {
+            "label": "水印间距",
+        },
+    },
+    "watermark_switch": {
+        "en": {
+            "label": "Watermark",
+            "value": "Not Add",
+            "choices": ["Not Add", "Add"],
+        },
+        "zh": {
+            "label": "水印",
+            "value": "不添加",
+            "choices": ["不添加", "添加"],
+        },
+    },
+    # 输出结果
+    "notification": {
+        "en": {
+            "label": "notification",
+            "face_error": "The number of faces is not equal to 1, please upload an image with a single face. If the actual number of faces is 1, it may be an issue with the accuracy of the detection model. Please switch to a different face detection model on the left or raise a Github Issue to notify the author.",
+        },
+        "zh": {
+            "label": "通知",
+            "face_error": "人脸数不等于1，请上传单人照片。如果实际人脸数为1，可能是检测模型的准确度问题，请切换左侧不同的人脸检测模型或提出Github Issue通知作者。",
+        },
+    },
+    "standard_photo": {
+        "en": {
+            "label": "Standard photo",
+        },
+        "zh": {
+            "label": "标准照",
+        },
+    },
+    "hd_photo": {
+        "en": {
+            "label": "HD photo",
+        },
+        "zh": {
+            "label": "高清照",
+        },
+    },
+    "layout_photo": {
+        "en": {
+            "label": "Layout photo",
+        },
+        "zh": {
+            "label": "六寸排版照",
+        },
+    },
+    "download": {
+        "en": {
+            "label": "Download the photo after adjusting the KB size",
+        },
+        "zh": {
+            "label": "下载调整 KB 大小后的照片",
+        },
+    },
+}

demo/processor.py

@@ -0,0 +1,286 @@
+import numpy as np
+from hivision import IDCreator
+from hivision.error import FaceError, APIError
+from hivision.utils import add_background, resize_image_to_kb
+from hivision.creator.layout_calculator import (
+    generate_layout_photo,
+    generate_layout_image,
+)
+from hivision.creator.choose_handler import choose_handler
+from demo.utils import add_watermark, range_check
+import gradio as gr
+import os
+import time
+from demo.locals import LOCALES
+
+
+class IDPhotoProcessor:
+    def process(
+        self,
+        input_image,
+        mode_option,
+        size_list_option,
+        color_option,
+        render_option,
+        image_kb_options,
+        custom_color_R,
+        custom_color_G,
+        custom_color_B,
+        custom_size_height,
+        custom_size_width,
+        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,
+        top_distance_max=0.12,
+        top_distance_min=0.10,
+    ):
+        top_distance_min = top_distance_max - 0.02
+
+        idphoto_json = {
+            "size_mode": mode_option,
+            "color_mode": color_option,
+            "render_mode": render_option,
+            "image_kb_mode": image_kb_options,
+            "custom_image_kb": None,
+        }
+
+        # 如果尺寸模式选择的是尺寸列表
+        if idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][0]:
+            idphoto_json["size"] = LOCALES["size_list"][language]["develop"][
+                size_list_option
+            ]
+        # 如果尺寸模式选择的是自定义尺寸
+        elif idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][2]:
+            id_height = int(custom_size_height)
+            id_width = int(custom_size_width)
+            if (
+                id_height < id_width
+                or min(id_height, id_width) < 100
+                or max(id_height, id_width) > 1800
+            ):
+                return [
+                    gr.update(value=None),  # img_output_standard
+                    gr.update(value=None),  # img_output_standard_hd
+                    None,  # img_output_layout (assuming it should be None or not updated)
+                    gr.update(  # notification
+                        value=LOCALES["size_mode"][language]["custom_size_eror"],
+                        visible=True,
+                    ),
+                    None,  # file_download (assuming it should be None or not updated)
+                ]
+
+            idphoto_json["size"] = (id_height, id_width)
+        else:
+            idphoto_json["size"] = (None, None)
+
+        # 如果颜色模式选择的是自定义底色
+        if idphoto_json["color_mode"] == LOCALES["bg_color"][language]["choices"][-1]:
+            idphoto_json["color_bgr"] = (
+                range_check(custom_color_R),
+                range_check(custom_color_G),
+                range_check(custom_color_B),
+            )
+        else:
+            hex_color = idphoto_json["color_bgr"] = LOCALES["bg_color"][language][
+                "develop"
+            ][color_option]
+            # 转为 RGB
+            idphoto_json["color_bgr"] = tuple(
+                int(hex_color[i : i + 2], 16) for i in (0, 2, 4)
+            )
+
+        # 如果输出 KB 大小选择的是自定义
+        if (
+            idphoto_json["image_kb_mode"]
+            == LOCALES["image_kb"][language]["choices"][-1]
+        ):
+            idphoto_json["custom_image_kb"] = custom_image_kb
+
+        creator = IDCreator()
+        choose_handler(creator, matting_model_option, face_detect_option)
+
+        # 是否只换底
+        change_bg_only = (
+            idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]
+        )
+
+        try:
+            result = creator(
+                input_image,
+                change_bg_only=change_bg_only,
+                size=idphoto_json["size"],
+                head_measure_ratio=head_measure_ratio,
+                head_top_range=(top_distance_max, top_distance_min),
+            )
+        except FaceError:
+            return [
+                gr.update(value=None),  # img_output_standard
+                gr.update(value=None),  # img_output_standard_hd
+                gr.update(visible=False),  # img_output_layout
+                gr.update(  # notification
+                    value=LOCALES["notification"][language]["face_error"],
+                    visible=True,
+                ),
+                None,  # file_download (assuming it should be None or have no update)
+            ]
+
+        except APIError as e:
+            return [
+                gr.update(value=None),  # img_output_standard
+                gr.update(value=None),  # img_output_standard_hd
+                gr.update(visible=False),  # img_output_layout
+                gr.update(  # notification
+                    value=LOCALES["notification"][language]["face_error"],
+                    visible=True,
+                ),
+                None,  # file_download (assuming it should be None or have no update)
+            ]
+
+        else:
+            (result_image_standard, result_image_hd, _, _) = result
+            if (
+                idphoto_json["render_mode"]
+                == LOCALES["render_mode"][language]["choices"][0]
+            ):
+                result_image_standard = np.uint8(
+                    add_background(result_image_standard, bgr=idphoto_json["color_bgr"])
+                )
+                result_image_hd = np.uint8(
+                    add_background(result_image_hd, bgr=idphoto_json["color_bgr"])
+                )
+            elif (
+                idphoto_json["render_mode"]
+                == LOCALES["render_mode"][language]["choices"][1]
+            ):
+                result_image_standard = np.uint8(
+                    add_background(
+                        result_image_standard,
+                        bgr=idphoto_json["color_bgr"],
+                        mode="updown_gradient",
+                    )
+                )
+                result_image_hd = np.uint8(
+                    add_background(
+                        result_image_hd,
+                        bgr=idphoto_json["color_bgr"],
+                        mode="updown_gradient",
+                    )
+                )
+            else:
+                result_image_standard = np.uint8(
+                    add_background(
+                        result_image_standard,
+                        bgr=idphoto_json["color_bgr"],
+                        mode="center_gradient",
+                    )
+                )
+                result_image_hd = np.uint8(
+                    add_background(
+                        result_image_hd,
+                        bgr=idphoto_json["color_bgr"],
+                        mode="center_gradient",
+                    )
+                )
+
+            # 如果只换底，就不生成排版照
+            if change_bg_only:
+                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],
+                )
+
+                if (
+                    watermark_option
+                    == LOCALES["watermark_switch"][language]["choices"][1]
+                ):
+                    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,
+                    )
+
+            # 如果添加水印
+            if watermark_option == LOCALES["watermark_switch"][language]["choices"][1]:
+                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,
+                )
+
+            # 如果输出 KB 大小选择的是自定义
+            if idphoto_json["custom_image_kb"]:
+                print("调整 kb 大小到", idphoto_json["custom_image_kb"], "kb")
+                output_image_path = f"{os.path.join(os.path.dirname(os.path.dirname(__file__)), 'demo/kb_output')}/{int(time.time())}.jpg"
+                resize_image_to_kb(
+                    result_image_standard,
+                    output_image_path,
+                    idphoto_json["custom_image_kb"],
+                )
+            else:
+                output_image_path = None
+
+            # 返回结果
+            if output_image_path:
+                return [
+                    result_image_standard,  # img_output_standard
+                    result_image_hd,  # img_output_standard_hd
+                    result_layout_image,  # img_output_layout
+                    gr.update(visible=False),  # notification
+                    gr.update(visible=True, value=output_image_path),  # file_download
+                ]
+            else:
+                return [
+                    result_image_standard,  # img_output_standard
+                    result_image_hd,  # img_output_standard_hd
+                    result_layout_image,  # img_output_layout
+                    gr.update(visible=False),  # notification
+                    gr.update(visible=False),  # file_download
+                ]

demo/ui.py

@@ -0,0 +1,493 @@
+import gradio as gr
+import os
+import pathlib
+from demo.locals import LOCALES
+from hivision.creator.choose_handler import FACE_DETECT_MODELS
+
+
+def load_description(fp):
+    with open(fp, "r", encoding="utf-8") as f:
+        content = f.read()
+    return content
+
+
+def create_ui(
+    processor, root_dir, human_matting_models: list, face_detect_models: list
+):
+    DEFAULT_LANG = "zh"
+    DEFAULT_HUMAN_MATTING_MODEL = "modnet_photographic_portrait_matting"
+    DEFAULT_FACE_DETECT_MODEL = "mtcnn"
+
+    if DEFAULT_HUMAN_MATTING_MODEL in human_matting_models:
+        human_matting_models.remove(DEFAULT_HUMAN_MATTING_MODEL)
+        human_matting_models.insert(0, DEFAULT_HUMAN_MATTING_MODEL)
+
+    css = """
+        #col-left {
+            margin: 0 auto;
+            max-width: 430px;
+        }
+        #col-mid {
+            margin: 0 auto;
+            max-width: 430px;
+        }
+        #col-right {
+            margin: 0 auto;
+            max-width: 430px;
+        }
+        #col-showcase {
+            margin: 0 auto;
+            max-width: 1100px;
+        }
+        #button {
+            color: blue;
+        }
+        """
+
+    demo = gr.Blocks(title="HivisionIDPhotos", css=css)
+
+    with demo:
+        gr.HTML(load_description(os.path.join(root_dir, "assets/title.md")))
+        with gr.Row():
+            # ------------ 左半边 UI ----------------
+            with gr.Column():
+                img_input = gr.Image(height=400)
+
+                with gr.Row():
+                    # 语言选择器
+                    language = ["zh", "en"]
+                    language_options = gr.Dropdown(
+                        choices=language,
+                        label="Language",
+                        value=DEFAULT_LANG,
+                    )
+
+                    face_detect_model_options = gr.Dropdown(
+                        choices=face_detect_models,
+                        label=LOCALES["face_model"][DEFAULT_LANG]["label"],
+                        value=DEFAULT_FACE_DETECT_MODEL,
+                    )
+
+                    matting_model_options = gr.Dropdown(
+                        choices=human_matting_models,
+                        label=LOCALES["matting_model"][DEFAULT_LANG]["label"],
+                        value=human_matting_models[0],
+                    )
+
+                with gr.Tab(
+                    LOCALES["key_param"][DEFAULT_LANG]["label"]
+                ) as key_parameter_tab:
+                    mode_options = gr.Radio(
+                        choices=LOCALES["size_mode"][DEFAULT_LANG]["choices"],
+                        label=LOCALES["size_mode"][DEFAULT_LANG]["label"],
+                        value=LOCALES["size_mode"][DEFAULT_LANG]["choices"][0],
+                    )
+
+                    with gr.Row(visible=True) as size_list_row:
+                        size_list_options = gr.Dropdown(
+                            choices=LOCALES["size_list"][DEFAULT_LANG]["choices"],
+                            label="预设尺寸",
+                            value=LOCALES["size_list"][DEFAULT_LANG]["choices"][0],
+                            elem_id="size_list",
+                        )
+
+                    with gr.Row(visible=False) as custom_size:
+                        custom_size_height = gr.Number(
+                            value=413, label="height", interactive=True
+                        )
+                        custom_size_width = gr.Number(
+                            value=295, label="width", interactive=True
+                        )
+
+                    color_options = gr.Radio(
+                        choices=LOCALES["bg_color"][DEFAULT_LANG]["choices"],
+                        label=LOCALES["bg_color"][DEFAULT_LANG]["label"],
+                        value=LOCALES["bg_color"][DEFAULT_LANG]["choices"][0],
+                    )
+
+                    with gr.Row(visible=False) as custom_color:
+                        custom_color_R = gr.Number(value=0, label="R", interactive=True)
+                        custom_color_G = gr.Number(value=0, label="G", interactive=True)
+                        custom_color_B = gr.Number(value=0, label="B", interactive=True)
+
+                    render_options = gr.Radio(
+                        choices=LOCALES["render_mode"][DEFAULT_LANG]["choices"],
+                        label=LOCALES["render_mode"][DEFAULT_LANG]["label"],
+                        value=LOCALES["render_mode"][DEFAULT_LANG]["choices"][0],
+                    )
+
+                with gr.Tab(
+                    LOCALES["advance_param"][DEFAULT_LANG]["label"]
+                ) as advance_parameter_tab:
+                    head_measure_ratio_option = gr.Slider(
+                        minimum=0.1,
+                        maximum=0.5,
+                        value=0.2,
+                        step=0.01,
+                        label=LOCALES["head_measure_ratio"][DEFAULT_LANG]["label"],
+                        interactive=True,
+                    )
+                    top_distance_option = gr.Slider(
+                        minimum=0.02,
+                        maximum=0.5,
+                        value=0.12,
+                        step=0.01,
+                        label=LOCALES["top_distance"][DEFAULT_LANG]["label"],
+                        interactive=True,
+                    )
+
+                    image_kb_options = gr.Radio(
+                        choices=LOCALES["image_kb"][DEFAULT_LANG]["choices"],
+                        label=LOCALES["image_kb"][DEFAULT_LANG]["label"],
+                        value=LOCALES["image_kb"][DEFAULT_LANG]["choices"][0],
+                    )
+
+                    with gr.Row(visible=False) as custom_image_kb:
+                        custom_image_kb_size = gr.Slider(
+                            minimum=10,
+                            maximum=1000,
+                            value=50,
+                            label=LOCALES["image_kb_size"][DEFAULT_LANG]["label"],
+                            interactive=True,
+                        )
+
+                with gr.Tab(
+                    LOCALES["watermark_tab"][DEFAULT_LANG]["label"]
+                ) as watermark_parameter_tab:
+                    watermark_options = gr.Radio(
+                        choices=LOCALES["watermark_switch"][DEFAULT_LANG]["choices"],
+                        label=LOCALES["watermark_switch"][DEFAULT_LANG]["label"],
+                        value=LOCALES["watermark_switch"][DEFAULT_LANG]["choices"][0],
+                    )
+
+                    with gr.Row():
+                        watermark_text_options = gr.Text(
+                            max_length=20,
+                            label=LOCALES["watermark_text"][DEFAULT_LANG]["label"],
+                            value=LOCALES["watermark_text"][DEFAULT_LANG]["value"],
+                            placeholder=LOCALES["watermark_text"][DEFAULT_LANG][
+                                "placeholder"
+                            ],
+                            interactive=False,
+                        )
+                        watermark_text_color = gr.ColorPicker(
+                            label=LOCALES["watermark_color"][DEFAULT_LANG]["label"],
+                            interactive=False,
+                            value="#FFFFFF",
+                        )
+
+                    watermark_text_size = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=20,
+                        label=LOCALES["watermark_size"][DEFAULT_LANG]["label"],
+                        interactive=False,
+                        step=1,
+                    )
+
+                    watermark_text_opacity = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        value=0.15,
+                        label=LOCALES["watermark_opacity"][DEFAULT_LANG]["label"],
+                        interactive=False,
+                        step=0.01,
+                    )
+
+                    watermark_text_angle = gr.Slider(
+                        minimum=0,
+                        maximum=360,
+                        value=30,
+                        label=LOCALES["watermark_angle"][DEFAULT_LANG]["label"],
+                        interactive=False,
+                        step=1,
+                    )
+
+                    watermark_text_space = gr.Slider(
+                        minimum=10,
+                        maximum=200,
+                        value=25,
+                        label=LOCALES["watermark_space"][DEFAULT_LANG]["label"],
+                        interactive=False,
+                        step=1,
+                    )
+
+                    def update_watermark_text_visibility(choice, language):
+                        return [
+                            gr.update(
+                                interactive=(
+                                    choice
+                                    == LOCALES["watermark_switch"][language]["choices"][
+                                        1
+                                    ]
+                                )
+                            )
+                        ] * 6
+
+                    watermark_options.change(
+                        fn=update_watermark_text_visibility,
+                        inputs=[watermark_options, language_options],
+                        outputs=[
+                            watermark_text_options,
+                            watermark_text_color,
+                            watermark_text_size,
+                            watermark_text_opacity,
+                            watermark_text_angle,
+                            watermark_text_space,
+                        ],
+                    )
+
+                img_but = gr.Button(LOCALES["button"][DEFAULT_LANG]["label"])
+
+                example_images = gr.Examples(
+                    inputs=[img_input],
+                    examples=[
+                        [path.as_posix()]
+                        for path in sorted(
+                            pathlib.Path(os.path.join(root_dir, "demo/images")).rglob(
+                                "*.jpg"
+                            )
+                        )
+                    ],
+                )
+
+            # ---------------- 右半边 UI ----------------
+            with gr.Column():
+                notification = gr.Text(
+                    label=LOCALES["notification"][DEFAULT_LANG]["label"], visible=False
+                )
+                with gr.Row():
+                    img_output_standard = gr.Image(
+                        label=LOCALES["standard_photo"][DEFAULT_LANG]["label"],
+                        height=350,
+                        format="jpeg",
+                    )
+                    img_output_standard_hd = gr.Image(
+                        label=LOCALES["hd_photo"][DEFAULT_LANG]["label"],
+                        height=350,
+                        format="jpeg",
+                    )
+                img_output_layout = gr.Image(
+                    label=LOCALES["layout_photo"][DEFAULT_LANG]["label"],
+                    height=350,
+                    format="jpeg",
+                )
+                file_download = gr.File(
+                    label=LOCALES["download"][DEFAULT_LANG]["label"], visible=False
+                )
+
+            # ---------------- 设置隐藏/显示组件 ----------------
+            def change_language(language):
+                return {
+                    face_detect_model_options: gr.update(
+                        label=LOCALES["face_model"][language]["label"]
+                    ),
+                    matting_model_options: gr.update(
+                        label=LOCALES["matting_model"][language]["label"]
+                    ),
+                    size_list_options: gr.update(
+                        label=LOCALES["size_list"][language]["label"],
+                        choices=LOCALES["size_list"][language]["choices"],
+                        value=LOCALES["size_list"][language]["choices"][0],
+                    ),
+                    mode_options: gr.update(
+                        label=LOCALES["size_mode"][language]["label"],
+                        choices=LOCALES["size_mode"][language]["choices"],
+                        value=LOCALES["size_mode"][language]["choices"][0],
+                    ),
+                    color_options: gr.update(
+                        label=LOCALES["bg_color"][language]["label"],
+                        choices=LOCALES["bg_color"][language]["choices"],
+                        value=LOCALES["bg_color"][language]["choices"][0],
+                    ),
+                    img_but: gr.update(value=LOCALES["button"][language]["label"]),
+                    render_options: gr.update(
+                        label=LOCALES["render_mode"][language]["label"],
+                        choices=LOCALES["render_mode"][language]["choices"],
+                        value=LOCALES["render_mode"][language]["choices"][0],
+                    ),
+                    image_kb_options: gr.update(
+                        label=LOCALES["image_kb_size"][language]["label"],
+                        choices=LOCALES["image_kb"][language]["choices"],
+                        value=LOCALES["image_kb"][language]["choices"][0],
+                    ),
+                    custom_image_kb_size: gr.update(
+                        label=LOCALES["image_kb"][language]["label"]
+                    ),
+                    notification: gr.update(
+                        label=LOCALES["notification"][language]["label"]
+                    ),
+                    img_output_standard: gr.update(
+                        label=LOCALES["standard_photo"][language]["label"]
+                    ),
+                    img_output_standard_hd: gr.update(
+                        label=LOCALES["hd_photo"][language]["label"]
+                    ),
+                    img_output_layout: gr.update(
+                        label=LOCALES["layout_photo"][language]["label"]
+                    ),
+                    file_download: gr.update(
+                        label=LOCALES["download"][language]["label"]
+                    ),
+                    head_measure_ratio_option: gr.update(
+                        label=LOCALES["head_measure_ratio"][language]["label"]
+                    ),
+                    top_distance_option: gr.update(
+                        label=LOCALES["top_distance"][language]["label"]
+                    ),
+                    key_parameter_tab: gr.update(
+                        label=LOCALES["key_param"][language]["label"]
+                    ),
+                    advance_parameter_tab: gr.update(
+                        label=LOCALES["advance_param"][language]["label"]
+                    ),
+                    watermark_parameter_tab: gr.update(
+                        label=LOCALES["watermark_tab"][language]["label"]
+                    ),
+                    watermark_text_options: gr.update(
+                        label=LOCALES["watermark_text"][language]["label"],
+                        placeholder=LOCALES["watermark_text"][language]["placeholder"],
+                    ),
+                    watermark_text_color: gr.update(
+                        label=LOCALES["watermark_color"][language]["label"]
+                    ),
+                    watermark_text_size: gr.update(
+                        label=LOCALES["watermark_size"][language]["label"]
+                    ),
+                    watermark_text_opacity: gr.update(
+                        label=LOCALES["watermark_opacity"][language]["label"]
+                    ),
+                    watermark_text_angle: gr.update(
+                        label=LOCALES["watermark_angle"][language]["label"]
+                    ),
+                    watermark_text_space: gr.update(
+                        label=LOCALES["watermark_space"][language]["label"]
+                    ),
+                    watermark_options: gr.update(
+                        label=LOCALES["watermark_switch"][language]["label"],
+                        choices=LOCALES["watermark_switch"][language]["choices"],
+                        value=LOCALES["watermark_switch"][language]["choices"][0],
+                    ),
+                }
+
+            def change_color(colors):
+                if colors == "自定义底色" or colors == "Custom Color":
+                    return {custom_color: gr.update(visible=True)}
+                else:
+                    return {custom_color: gr.update(visible=False)}
+
+            def change_size_mode(size_option_item):
+                if (
+                    size_option_item == "自定义尺寸"
+                    or size_option_item == "Custom Size"
+                ):
+                    return {
+                        custom_size: gr.update(visible=True),
+                        size_list_row: gr.update(visible=False),
+                    }
+                elif (
+                    size_option_item == "只换底"
+                    or size_option_item == "Only Change Background"
+                ):
+                    return {
+                        custom_size: gr.update(visible=False),
+                        size_list_row: gr.update(visible=False),
+                    }
+                else:
+                    return {
+                        custom_size: gr.update(visible=False),
+                        size_list_row: gr.update(visible=True),
+                    }
+
+            def change_image_kb(image_kb_option):
+                if image_kb_option == "自定义" or image_kb_option == "Custom":
+                    return {custom_image_kb: gr.update(visible=True)}
+                else:
+                    return {custom_image_kb: gr.update(visible=False)}
+
+            # ---------------- 绑定事件 ----------------
+            language_options.input(
+                change_language,
+                inputs=[language_options],
+                outputs=[
+                    size_list_options,
+                    mode_options,
+                    color_options,
+                    img_but,
+                    render_options,
+                    image_kb_options,
+                    matting_model_options,
+                    face_detect_model_options,
+                    custom_image_kb_size,
+                    notification,
+                    img_output_standard,
+                    img_output_standard_hd,
+                    img_output_layout,
+                    file_download,
+                    head_measure_ratio_option,
+                    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,
+                ],
+            )
+
+            color_options.input(
+                change_color, inputs=[color_options], outputs=[custom_color]
+            )
+
+            mode_options.input(
+                change_size_mode,
+                inputs=[mode_options],
+                outputs=[custom_size, size_list_row],
+            )
+
+            image_kb_options.input(
+                change_image_kb, inputs=[image_kb_options], outputs=[custom_image_kb]
+            )
+
+            img_but.click(
+                processor.process,
+                inputs=[
+                    img_input,
+                    mode_options,
+                    size_list_options,
+                    color_options,
+                    render_options,
+                    image_kb_options,
+                    custom_color_R,
+                    custom_color_G,
+                    custom_color_B,
+                    custom_size_height,
+                    custom_size_width,
+                    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,
+                ],
+                outputs=[
+                    img_output_standard,
+                    img_output_standard_hd,
+                    img_output_layout,
+                    notification,
+                    file_download,
+                ],
+            )
+
+    return demo

demo/utils.py

@@ -1,4 +1,7 @@
 import csv
+import numpy as np
+from PIL import Image
+from hivision.plugin.watermark import Watermarker, WatermarkerStyles
 
 
 def csv_to_size_list(csv_file: str) -> dict:
@@ -17,3 +20,42 @@ def csv_to_size_list(csv_file: str) -> dict:
             size_list_dict[size_name_add_size] = (int(h), int(w))
 
     return size_list_dict
+
+
+def csv_to_color_list(csv_file: str) -> dict:
+    # 初始化一个空字典
+    color_list_dict = {}
+
+    # 打开 CSV 文件并读取数据
+    with open(csv_file, mode="r", encoding="utf-8") as file:
+        reader = csv.reader(file)
+        # 跳过表头
+        next(reader)
+        # 读取数据并填充字典
+        for row in reader:
+            color_name, hex_code = row
+            color_list_dict[color_name] = hex_code
+
+    return color_list_dict
+
+
+def range_check(value, min_value=0, max_value=255):
+    value = int(value)
+    return max(min_value, min(value, max_value))
+
+
+def add_watermark(
+    image, text, size=50, opacity=0.5, angle=45, color="#8B8B1B", space=75
+):
+    image = Image.fromarray(image)
+    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"))

hivision/creator/choose_handler.py

@@ -2,6 +2,16 @@ from hivision.creator.human_matting import *
 from hivision.creator.face_detector import *
 
 
+HUMAN_MATTING_MODELS = [
+    "modnet_photographic_portrait_matting",
+    "birefnet-v1-lite",
+    "hivision_modnet",
+    "rmbg-1.4",
+]
+
+FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn", "retinaface-resnet50"]
+
+
 def choose_handler(creator, matting_model_option=None, face_detect_option=None):
     if matting_model_option == "modnet_photographic_portrait_matting":
         creator.matting_handler = extract_human_modnet_photographic_portrait_matting
@@ -9,8 +19,6 @@ def choose_handler(creator, matting_model_option=None, face_detect_option=None):
         creator.matting_handler = extract_human_mnn_modnet
     elif matting_model_option == "rmbg-1.4":
         creator.matting_handler = extract_human_rmbg
-    # elif matting_model_option == "birefnet-portrait":
-    #     creator.matting_handler = extract_human_birefnet_portrait
     elif matting_model_option == "birefnet-v1-lite":
         creator.matting_handler = extract_human_birefnet_lite
     else:
@@ -18,8 +26,10 @@ def choose_handler(creator, matting_model_option=None, face_detect_option=None):
 
     if (
         face_detect_option == "face_plusplus"
-        or face_detect_option == "face++ (联网API)"
+        or face_detect_option == "face++ (联网Online API)"
     ):
         creator.detection_handler = detect_face_face_plusplus
+    elif face_detect_option == "retinaface-resnet50":
+        creator.detection_handler = detect_face_retinaface
     else:
         creator.detection_handler = detect_face_mtcnn

hivision/creator/face_detector.py

@@ -7,16 +7,24 @@ r"""
 @Description:
     人脸检测器
 """
-from mtcnnruntime import MTCNN
+try:
+    from mtcnnruntime import MTCNN
+except ImportError:
+    raise ImportError(
+        "Please install mtcnn-runtime by running `pip install mtcnn-runtime`"
+    )
 from .context import Context
 from hivision.error import FaceError, APIError
 from hivision.utils import resize_image_to_kb_base64
+from hivision.creator.retinaface import retinaface_detect_faces
 import requests
 import cv2
 import os
 
 
 mtcnn = None
+base_dir = os.path.dirname(os.path.abspath(__file__))
+RETINAFCE_SESS = None
 
 
 def detect_face_mtcnn(ctx: Context, scale: int = 2):
@@ -124,3 +132,45 @@ def detect_face_face_plusplus(ctx: Context):
             f"Face++ Status code {status_code} Request entity too large: The image exceeds the 2MB limit.",
             status_code,
         )
+
+
+def detect_face_retinaface(ctx: Context):
+    """
+    基于RetinaFace模型的人脸检测处理器，只进行人脸数量的检测
+    :param ctx: 上下文，此时已获取到原始图和抠图结果，但是我们只需要原始图
+    :raise FaceError: 人脸检测错误，多个人脸或者没有人脸
+    """
+    from time import time
+
+    global RETINAFCE_SESS
+
+    if RETINAFCE_SESS is None:
+        print("首次加载RetinaFace模型...")
+        # 计算用时
+        tic = time()
+        faces_dets, sess = retinaface_detect_faces(
+            ctx.origin_image,
+            os.path.join(base_dir, "retinaface/weights/retinaface-resnet50.onnx"),
+            sess=None,
+        )
+        RETINAFCE_SESS = sess
+        print("首次RetinaFace模型推理用时: {:.4f}s".format(time() - tic))
+    else:
+        tic = time()
+        faces_dets, _ = retinaface_detect_faces(
+            ctx.origin_image,
+            os.path.join(base_dir, "retinaface/weights/retinaface-resnet50.onnx"),
+            sess=RETINAFCE_SESS,
+        )
+        print("二次RetinaFace模型推理用时: {:.4f}s".format(time() - tic))
+
+    faces_num = len(faces_dets)
+    if faces_num != 1:
+        raise FaceError("Expected 1 face, but got {}".format(faces_num), faces_num)
+    face_det = faces_dets[0]
+    ctx.face = (
+        face_det[0],
+        face_det[1],
+        face_det[2] - face_det[0] + 1,
+        face_det[3] - face_det[1] + 1,
+    )

hivision/creator/retinaface/__init__.py

@@ -0,0 +1 @@
+from .inference import retinaface_detect_faces

hivision/creator/retinaface/box_utils.py

@@ -0,0 +1,57 @@
+import numpy as np
+
+
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = None
+
+    boxes = np.concatenate(
+        (
+            priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+            priors[:, 2:] * np.exp(loc[:, 2:] * variances[1]),
+        ),
+        axis=1,
+    )
+
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes
+
+
+def decode_landm(pre, priors, variances):
+    """Decode landm from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        pre (tensor): landm predictions for loc layers,
+            Shape: [num_priors,10]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded landm predictions
+    """
+    landms = None
+
+    landms = np.concatenate(
+        (
+            priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+            priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+        ),
+        axis=1,
+    )
+
+    return landms

hivision/creator/retinaface/inference.py

@@ -0,0 +1,190 @@
+import numpy as np
+import cv2
+import onnxruntime as ort
+from hivision.creator.retinaface.box_utils import decode, decode_landm
+from hivision.creator.retinaface.prior_box import PriorBox
+import argparse
+
+
+def py_cpu_nms(dets, thresh):
+    """Pure Python NMS baseline."""
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+    scores = dets[:, 4]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+parser = argparse.ArgumentParser(description="Retinaface")
+
+parser.add_argument(
+    "--network", default="resnet50", help="Backbone network mobile0.25 or resnet50"
+)
+parser.add_argument(
+    "--cpu", action="store_true", default=False, help="Use cpu inference"
+)
+parser.add_argument(
+    "--confidence_threshold", default=0.8, type=float, help="confidence_threshold"
+)
+parser.add_argument("--top_k", default=5000, type=int, help="top_k")
+parser.add_argument("--nms_threshold", default=0.2, type=float, help="nms_threshold")
+parser.add_argument("--keep_top_k", default=750, type=int, help="keep_top_k")
+parser.add_argument(
+    "-s",
+    "--save_image",
+    action="store_true",
+    default=True,
+    help="show detection results",
+)
+parser.add_argument(
+    "--vis_thres", default=0.6, type=float, help="visualization_threshold"
+)
+args = parser.parse_args()
+
+
+def load_model_ort(model_path):
+    ort_session = ort.InferenceSession(model_path)
+    return ort_session
+
+
+def retinaface_detect_faces(image, model_path: str, sess=None):
+    cfg = {
+        "name": "Resnet50",
+        "min_sizes": [[16, 32], [64, 128], [256, 512]],
+        "steps": [8, 16, 32],
+        "variance": [0.1, 0.2],
+        "clip": False,
+        "loc_weight": 2.0,
+        "gpu_train": True,
+        "batch_size": 24,
+        "ngpu": 4,
+        "epoch": 100,
+        "decay1": 70,
+        "decay2": 90,
+        "image_size": 840,
+        "pretrain": True,
+        "return_layers": {"layer2": 1, "layer3": 2, "layer4": 3},
+        "in_channel": 256,
+        "out_channel": 256,
+    }
+
+    # Load ONNX model
+    if sess is None:
+        retinaface = load_model_ort(model_path)
+    else:
+        retinaface = sess
+
+    resize = 1
+
+    # Read and preprocess the image
+    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    img = np.float32(img_rgb)
+
+    im_height, im_width, _ = img.shape
+    scale = np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
+    img -= (104, 117, 123)
+    img = img.transpose(2, 0, 1)
+    img = np.expand_dims(img, axis=0)
+
+    # Run the model
+    inputs = {"input": img}
+    loc, conf, landms = retinaface.run(None, inputs)
+
+    # tic = time.time()
+    priorbox = PriorBox(cfg, image_size=(im_height, im_width))
+    priors = priorbox.forward()
+
+    prior_data = priors
+
+    boxes = decode(np.squeeze(loc, axis=0), prior_data, cfg["variance"])
+    boxes = boxes * scale / resize
+    scores = np.squeeze(conf, axis=0)[:, 1]
+
+    landms = decode_landm(np.squeeze(landms.data, axis=0), prior_data, cfg["variance"])
+
+    scale1 = np.array(
+        [
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+            img.shape[3],
+            img.shape[2],
+        ]
+    )
+    landms = landms * scale1 / resize
+
+    # ignore low scores
+    inds = np.where(scores > args.confidence_threshold)[0]
+    boxes = boxes[inds]
+    landms = landms[inds]
+    scores = scores[inds]
+
+    # keep top-K before NMS
+    order = scores.argsort()[::-1][: args.top_k]
+    boxes = boxes[order]
+    landms = landms[order]
+    scores = scores[order]
+
+    # do NMS
+    dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
+    keep = py_cpu_nms(dets, args.nms_threshold)
+    # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
+    dets = dets[keep, :]
+    landms = landms[keep]
+
+    # keep top-K faster NMS
+    dets = dets[: args.keep_top_k, :]
+    landms = landms[: args.keep_top_k, :]
+
+    dets = np.concatenate((dets, landms), axis=1)
+    # print("post processing time: {:.4f}s".format(time.time() - tic))
+
+    return dets, retinaface
+
+
+if __name__ == "__main__":
+    import gradio as gr
+
+    # Create Gradio interface
+    iface = gr.Interface(
+        fn=retinaface_detect_faces,
+        inputs=[
+            gr.Image(
+                type="numpy", label="上传图片", height=400
+            ),  # Set the height to 400
+            gr.Textbox(value="./FaceDetector.onnx", label="ONNX模型路径"),
+        ],
+        outputs=gr.Number(label="检测到的人脸数量"),
+        title="人脸检测",
+        description="上传图片并提供ONNX模型路径以检测人脸数量。",
+    )
+
+    # Launch the Gradio app
+    iface.launch()

hivision/creator/retinaface/prior_box.py

@@ -0,0 +1,41 @@
+from itertools import product as product
+import numpy as np
+from math import ceil
+
+
+class PriorBox(object):
+    def __init__(self, cfg, image_size=None):
+        super(PriorBox, self).__init__()
+        self.min_sizes = cfg["min_sizes"]
+        self.steps = cfg["steps"]
+        self.clip = cfg["clip"]
+        self.image_size = image_size
+        self.feature_maps = [
+            [ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)]
+            for step in self.steps
+        ]
+        self.name = "s"
+
+    def forward(self):
+        anchors = []
+        for k, f in enumerate(self.feature_maps):
+            min_sizes = self.min_sizes[k]
+            for i, j in product(range(f[0]), range(f[1])):
+                for min_size in min_sizes:
+                    s_kx = min_size / self.image_size[1]
+                    s_ky = min_size / self.image_size[0]
+                    dense_cx = [
+                        x * self.steps[k] / self.image_size[1] for x in [j + 0.5]
+                    ]
+                    dense_cy = [
+                        y * self.steps[k] / self.image_size[0] for y in [i + 0.5]
+                    ]
+                    for cy, cx in product(dense_cy, dense_cx):
+                        anchors += [cx, cy, s_kx, s_ky]
+
+        output = np.array(anchors).reshape(-1, 4)
+
+        if self.clip:
+            output = np.clip(output, 0, 1)
+
+        return output

hivision/creator/retinaface/weights/retinaface-resnet50.onnx

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