demo/processor.py

import numpy as np
from hivision import IDCreator
from hivision.error import FaceError, APIError
from hivision.utils import (
    add_background,
    add_background_with_image,
    resize_image_to_kb,
    add_watermark,
    save_image_dpi_to_bytes,
)
from hivision.creator.layout_calculator import (
    generate_layout_photo,
    generate_layout_image,
)
from hivision.creator.choose_handler import choose_handler
from hivision.plugin.template.template_calculator import generte_template_photo
from demo.utils import range_check
import gradio as gr
import os
import cv2
import time
from demo.locales import LOCALES


base_path = os.path.dirname(os.path.abspath(__file__))

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_color_hex_value,
        custom_size_height,
        custom_size_width,
        custom_size_height_mm,
        custom_size_width_mm,
        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,
        whitening_strength=0,
        image_dpi_option=False,
        custom_image_dpi=None,
        brightness_strength=0,
        contrast_strength=0,
        sharpen_strength=0,
        saturation_strength=0,
        face_alignment_option=False,
    ):
        # 初始化参数
        top_distance_min = top_distance_max - 0.02
        # 得到render_option在LOCALES["render_mode"][language]["choices"]中的索引
        render_option_index = LOCALES["render_mode"][language]["choices"].index(
            render_option
        )
        idphoto_json = self._initialize_idphoto_json(
            mode_option, color_option, render_option_index, image_kb_options
        )

        # 处理尺寸模式
        size_result = self._process_size_mode(
            idphoto_json,
            language,
            size_list_option,
            custom_size_height,
            custom_size_width,
            custom_size_height_mm,
            custom_size_width_mm,
        )
        if isinstance(size_result, list):
            return size_result  # 返回错误信息

        # 处理颜色模式
        self._process_color_mode(
            idphoto_json,
            language,
            color_option,
            custom_color_R,
            custom_color_G,
            custom_color_B,
            custom_color_hex_value,
        )

        # 如果设置了自定义KB大小
        if (
            idphoto_json["image_kb_mode"]
            == LOCALES["image_kb"][language]["choices"][-1]
        ):
            idphoto_json["custom_image_kb"] = custom_image_kb

        # 如果设置了自定义DPI大小
        if image_dpi_option == LOCALES["image_dpi"][language]["choices"][-1]:
            idphoto_json["custom_image_dpi"] = custom_image_dpi

        # 创建IDCreator实例并设置处理器
        creator = IDCreator()
        choose_handler(creator, matting_model_option, face_detect_option)

        # 生成证件照
        try:
            result = self._generate_id_photo(
                creator,
                input_image,
                idphoto_json,
                language,
                head_measure_ratio,
                top_distance_max,
                top_distance_min,
                whitening_strength,
                brightness_strength,
                contrast_strength,
                sharpen_strength,
                saturation_strength,
                face_alignment_option,
            )
        except (FaceError, APIError):
            return self._handle_photo_generation_error(language)

        # 后处理生成的照片
        return self._process_generated_photo(
            result,
            idphoto_json,
            language,
            watermark_option,
            watermark_text,
            watermark_text_size,
            watermark_text_opacity,
            watermark_text_angle,
            watermark_text_space,
            watermark_text_color,
        )

    # 初始化idphoto_json字典
    def _initialize_idphoto_json(
        self,
        mode_option,
        color_option,
        render_option,
        image_kb_options,
    ):
        """初始化idphoto_json字典"""
        return {
            "size_mode": mode_option,
            "color_mode": color_option,
            "render_mode": render_option,
            "image_kb_mode": image_kb_options,
            "custom_image_kb": None,
            "custom_image_dpi": None,
        }

    # 处理尺寸模式
    def _process_size_mode(
        self,
        idphoto_json,
        language,
        size_list_option,
        custom_size_height,
        custom_size_width,
        custom_size_height_mm,
        custom_size_width_mm,
    ):
        """处理尺寸模式"""
        # 如果选择了尺寸列表
        if idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][0]:
            idphoto_json["size"] = LOCALES["size_list"][language]["develop"][
                size_list_option
            ]
        # 如果选择了自定义尺寸(px或mm)
        elif (
            idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][2]
            or idphoto_json["size_mode"] == LOCALES["size_mode"][language]["choices"][3]
        ):
            # 如果选择了自定义尺寸(px)
            if (
                idphoto_json["size_mode"]
                == LOCALES["size_mode"][language]["choices"][2]
            ):
                id_height, id_width = int(custom_size_height), int(custom_size_width)
            # 如果选择了自定义尺寸(mm)
            else:
                # 将mm转换为px
                id_height = int(custom_size_height_mm / 25.4 * 300)
                id_width = int(custom_size_width_mm / 25.4 * 300)
            # 检查尺寸像素是否在100到1800之间
            if (
                id_height < id_width
                or min(id_height, id_width) < 100
                or max(id_height, id_width) > 1800
            ):
                return self._create_error_response(language)
            idphoto_json["size"] = (id_height, id_width)
        # 如果选择了只换底
        else:
            idphoto_json["size"] = (None, None)

    # 处理颜色模式
    def _process_color_mode(
        self,
        idphoto_json,
        language,
        color_option,
        custom_color_R,
        custom_color_G,
        custom_color_B,
        custom_color_hex_value,
    ):
        """处理颜色模式"""
        # 如果选择了自定义颜色BGR
        if idphoto_json["color_mode"] == LOCALES["bg_color"][language]["choices"][-2]:
            idphoto_json["color_bgr"] = tuple(
                map(range_check, [custom_color_R, custom_color_G, custom_color_B])
            )
        # 如果选择了自定义颜色HEX
        elif idphoto_json["color_mode"] == LOCALES["bg_color"][language]["choices"][-1]:
            hex_color = custom_color_hex_value
            # 将十六进制颜色转换为RGB颜色，如果长度为6，则直接转换，如果长度为7，则去掉#号再转换
            if len(hex_color) == 6:
                idphoto_json["color_bgr"] = tuple(
                    int(hex_color[i : i + 2], 16) for i in (0, 2, 4)
                )
            elif len(hex_color) == 7:
                hex_color = hex_color[1:]
                idphoto_json["color_bgr"] = tuple(
                    int(hex_color[i : i + 2], 16) for i in (0, 2, 4)
                )
            else:
                raise ValueError(
                    "Invalid hex color. You can only use 6 or 7 characters. For example: #FFFFFF or FFFFFF"
                )
        # 如果选择了美式证件照
        elif idphoto_json["color_mode"] == LOCALES["bg_color"][language]["choices"][-3]:
            idphoto_json["color_bgr"] = (255, 255, 255)
        else:
            hex_color = LOCALES["bg_color"][language]["develop"][color_option]
            idphoto_json["color_bgr"] = tuple(
                int(hex_color[i : i + 2], 16) for i in (0, 2, 4)
            )

    # 生成证件照
    def _generate_id_photo(
        self,
        creator: IDCreator,
        input_image,
        idphoto_json,
        language,
        head_measure_ratio,
        top_distance_max,
        top_distance_min,
        whitening_strength,
        brightness_strength,
        contrast_strength,
        sharpen_strength,
        saturation_strength,
        face_alignment_option,
    ):
        """生成证件照"""
        change_bg_only = (
            idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]
        )
        return 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),
            whitening_strength=whitening_strength,
            brightness_strength=brightness_strength,
            contrast_strength=contrast_strength,
            sharpen_strength=sharpen_strength,
            saturation_strength=saturation_strength,
            face_alignment=face_alignment_option,
        )

    # 处理照片生成错误
    def _handle_photo_generation_error(self, language):
        """处理照片生成错误"""
        return [gr.update(value=None) for _ in range(4)] + [
            gr.update(visible=False),
            gr.update(
                value=LOCALES["notification"][language]["face_error"], visible=True
            ),
            None,
        ]

    # 处理生成的照片
    def _process_generated_photo(
        self,
        result,
        idphoto_json,
        language,
        watermark_option,
        watermark_text,
        watermark_text_size,
        watermark_text_opacity,
        watermark_text_angle,
        watermark_text_space,
        watermark_text_color,
    ):
        """处理生成的照片"""
        result_image_standard, result_image_hd, _, _, _, _ = result
        result_image_standard_png = np.uint8(result_image_standard)
        result_image_hd_png = np.uint8(result_image_hd)

        # 渲染背景
        result_image_standard, result_image_hd = self._render_background(
            result_image_standard, result_image_hd, idphoto_json, language
        )

        # 添加水印
        if watermark_option == LOCALES["watermark_switch"][language]["choices"][1]:
            result_image_standard, result_image_hd = self._add_watermark(
                result_image_standard,
                result_image_hd,
                watermark_text,
                watermark_text_size,
                watermark_text_opacity,
                watermark_text_angle,
                watermark_text_space,
                watermark_text_color,
            )
        
        # 生成排版照片
        result_image_layout, result_image_layout_visible = self._generate_image_layout(
            idphoto_json,
            result_image_standard,
            language,
        )
        
        # 生成模板照片
        result_image_template, result_image_template_visible = self._generate_image_template(
            idphoto_json,
            result_image_hd,
            language,
        )

        # 调整图片大小
        output_image_path_dict = self._resize_image_if_needed(
            result_image_standard,
            result_image_hd,
            result_image_layout,
            idphoto_json,
        )

        # 如果output_image_path_dict为None，即没有设置KB和DPI
        if output_image_path_dict is None:
            return self._create_response(
                result_image_standard,
                result_image_hd,
                result_image_standard_png,
                result_image_hd_png,
                gr.update(value=result_image_layout, visible=result_image_layout_visible),
                gr.update(value=result_image_template, visible=result_image_template_visible),
                gr.update(visible = result_image_template_visible),
            )
        # 如果output_image_path_dict不为None，即设置了KB和DPI
        else:
            if output_image_path_dict["layout"]["processed"]:
                result_image_layout = output_image_path_dict["layout"]["path"]
            return self._create_response(
                (
                    output_image_path_dict["standard"]["path"]
                    if output_image_path_dict["standard"]["processed"]
                    else result_image_standard
                ),
                (
                    output_image_path_dict["hd"]["path"]
                    if output_image_path_dict["hd"]["processed"]
                    else result_image_hd    
                ),
                result_image_standard_png,
                result_image_hd_png,
                gr.update(value=result_image_layout, visible=result_image_layout_visible),
                gr.update(value=result_image_template, visible=result_image_template_visible),
                gr.update(visible = result_image_template_visible),
            )

    # 渲染背景
    def _render_background(self, result_image_standard, result_image_hd, idphoto_json, language):
        """渲染背景"""
        render_modes = {0: "pure_color", 1: "updown_gradient", 2: "center_gradient"}
        render_mode = render_modes[idphoto_json["render_mode"]]

        if idphoto_json["color_mode"] != LOCALES["bg_color"][language]["choices"][-3]:
            result_image_standard = np.uint8(
                add_background(
                    result_image_standard, bgr=idphoto_json["color_bgr"], mode=render_mode
                )
            )
            result_image_hd = np.uint8(
                add_background(
                    result_image_hd, bgr=idphoto_json["color_bgr"], mode=render_mode
                )
            )
        # 如果选择了美式证件照
        else:
            result_image_standard = np.uint8(
                add_background_with_image(
                    result_image_standard, 
                    background_image=cv2.imread(os.path.join(base_path, "assets", "american-style.png"))
                )
            )
            result_image_hd = np.uint8(
                add_background_with_image(
                    result_image_hd, 
                    background_image=cv2.imread(os.path.join(base_path, "assets", "american-style.png"))
                )
            )
        return result_image_standard, result_image_hd

    # 生成排版照片
    def _generate_image_layout(
        self,
        idphoto_json,
        result_image_standard,
        language,
    ):
        """生成排版照片"""
        # 如果选择了只换底，则不生成排版照片
        if idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]:
            return None, False

        typography_arr, typography_rotate = generate_layout_photo(
            input_height=idphoto_json["size"][0],
            input_width=idphoto_json["size"][1],
        )
        
        result_image_layout = generate_layout_image(
            result_image_standard,
            typography_arr,
            typography_rotate,
            height=idphoto_json["size"][0],
            width=idphoto_json["size"][1],
        )

        return result_image_layout, True
    
    # 生成模板照片
    def _generate_image_template(
        self,
        idphoto_json,
        result_image_hd,
        language,
    ):
        # 如果选择了只换底，则不生成模板照片
        if idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]:
            return None, False
        
        TEMPLATE_NAME_LIST = ["template_1", "template_2"]
        """生成模板照片"""
        result_image_template_list = []
        for template_name in TEMPLATE_NAME_LIST:
            result_image_template = generte_template_photo(
                template_name=template_name,
                input_image=result_image_hd,
            )
            result_image_template_list.append(result_image_template)
        return result_image_template_list, True

    # 添加水印
    def _add_watermark(
        self,
        result_image_standard,
        result_image_hd,
        watermark_text,
        watermark_text_size,
        watermark_text_opacity,
        watermark_text_angle,
        watermark_text_space,
        watermark_text_color,
    ):
        """添加水印"""
        watermark_params = {
            "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_standard = add_watermark(
            image=result_image_standard, **watermark_params
        )
        result_image_hd = add_watermark(image=result_image_hd, **watermark_params)
        return result_image_standard, result_image_hd

    def _resize_image_if_needed(
        self,
        result_image_standard,
        result_image_hd,
        result_image_layout,
        idphoto_json,
        format="png",
    ):
        """如果需要，调整图片大小"""
        # 设置输出路径
        base_path = os.path.join(
            os.path.dirname(os.path.dirname(__file__)), "demo/kb_output"
        )
        timestamp = int(time.time())
        output_paths = {
            "standard": {
                "path": f"{base_path}/{timestamp}_standard",
                "processed": False,
            },
            "hd": {"path": f"{base_path}/{timestamp}_hd", "processed": False},
            "layout": {"path": f"{base_path}/{timestamp}_layout", "processed": False},
        }

        # 获取自定义的KB和DPI值
        custom_kb = idphoto_json.get("custom_image_kb")
        custom_dpi = idphoto_json.get("custom_image_dpi", 300)

        # 处理同时有自定义KB和DPI的情况
        if custom_kb and custom_dpi:
            # 为所有输出路径添加DPI信息
            for key in output_paths:
                output_paths[key]["path"] += f"_{custom_dpi}dpi.{format}"
            # 为标准图像添加KB信息
            output_paths["standard"]["path"] = output_paths["standard"]["path"].replace(
                f".{format}", f"_{custom_kb}kb.{format}"
            )

            # 调整标准图像大小并保存
            resize_image_to_kb(
                result_image_standard,
                output_paths["standard"]["path"],
                custom_kb,
                dpi=custom_dpi,
            )
            output_paths["standard"]["processed"] = True
            # 保存高清图像和排版图像
            save_image_dpi_to_bytes(
                result_image_hd, output_paths["hd"]["path"], dpi=custom_dpi
            )
            output_paths["hd"]["processed"] = True
            if result_image_layout is not None:
                save_image_dpi_to_bytes(
                    result_image_layout, output_paths["layout"]["path"], dpi=custom_dpi
                )
                output_paths["layout"]["processed"] = True

            return output_paths

        # 只有自定义DPI的情况
        elif custom_dpi:
            for key in output_paths:
                output_paths[key]["path"] += f"_{custom_dpi}dpi.{format}"
                # 保存所有图像，使用自定义DPI
                if key == "layout" and result_image_layout is None:
                    pass
                else:
                    save_image_dpi_to_bytes(
                        locals()[f"result_image_{key}"],
                        output_paths[key]["path"],
                        dpi=custom_dpi,
                    )
                    output_paths[key]["processed"] = True

            return output_paths

        # 只有自定义KB的情况
        elif custom_kb:
            output_paths["standard"]["path"] += f"_{custom_kb}kb.{format}"
            # 只调整标准图像大小并保存
            resize_image_to_kb(
                result_image_standard,
                output_paths["standard"]["path"],
                custom_kb,
                dpi=300,
            )
            output_paths["standard"]["processed"] = True

            return output_paths

        # 如果没有自定义设置，返回None
        return None

    def _create_response(
        self,
        result_image_standard,
        result_image_hd,
        result_image_standard_png,
        result_image_hd_png,
        result_layout_image_gr,
        result_image_template_gr,
        result_image_template_accordion_gr,
    ):
        """创建响应"""
        response = [
            result_image_standard,
            result_image_hd,
            result_image_standard_png,
            result_image_hd_png,
            result_layout_image_gr,
            result_image_template_gr,
            result_image_template_accordion_gr,
            gr.update(visible=False),
        ]

        return response

    def _create_error_response(self, language):
        """创建错误响应"""
        return [gr.update(value=None) for _ in range(4)] + [
            None,
            gr.update(
                value=LOCALES["size_mode"][language]["custom_size_eror"], visible=True
            ),
            None,
        ]