demo/processor.py

import numpy as np
from hivision import IDCreator
from hivision.error import FaceError, APIError
from hivision.utils import (
    add_background,
    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 demo.utils import 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,
        whitening_strength=0,
        image_dpi_option=False,
        custom_image_dpi=None,
        brightness_strength=0,
        contrast_strength=0,
        sharpen_strength=0,
        saturation_strength=0,
    ):
        # 初始化参数
        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,
        )
        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,
        )

        # 如果设置了自定义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,
            )
        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,
        )

    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,
    ):
        """处理尺寸模式"""
        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, id_width = int(custom_size_height), int(custom_size_width)
            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,
    ):
        """处理颜色模式"""
        if idphoto_json["color_mode"] == LOCALES["bg_color"][language]["choices"][-1]:
            idphoto_json["color_bgr"] = tuple(
                map(range_check, [custom_color_R, custom_color_G, custom_color_B])
            )
        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,
    ):
        """生成证件照"""
        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,
        )

    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
        )

        # 生成排版照片
        result_layout_image = self._generate_layout_image(
            idphoto_json,
            result_image_standard,
            language,
            watermark_option,
            watermark_text,
            watermark_text_size,
            watermark_text_opacity,
            watermark_text_angle,
            watermark_text_space,
            watermark_text_color,
        )

        # 添加水印
        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,
            )

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

        return self._create_response(
            result_image_standard,
            result_image_hd,
            result_image_standard_png,
            result_image_hd_png,
            result_layout_image,
            output_image_path,
        )

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

        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
            )
        )

        return result_image_standard, result_image_hd

    def _generate_layout_image(
        self,
        idphoto_json,
        result_image_standard,
        language,
        watermark_option,
        watermark_text,
        watermark_text_size,
        watermark_text_opacity,
        watermark_text_angle,
        watermark_text_space,
        watermark_text_color,
    ):
        """生成排版照片"""
        if idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]:
            return gr.update(visible=False)

        typography_arr, typography_rotate = generate_layout_photo(
            input_height=idphoto_json["size"][0],
            input_width=idphoto_json["size"][1],
        )

        image = result_image_standard
        if watermark_option == LOCALES["watermark_switch"][language]["choices"][1]:
            image = add_watermark(
                image=image,
                text=watermark_text,
                size=watermark_text_size,
                opacity=watermark_text_opacity,
                angle=watermark_text_angle,
                space=watermark_text_space,
                color=watermark_text_color,
            )

        return gr.update(
            value=generate_layout_image(
                image,
                typography_arr,
                typography_rotate,
                height=idphoto_json["size"][0],
                width=idphoto_json["size"][1],
            ),
            visible=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, idphoto_json):
        """如果需要，调整图片大小"""
        output_image_path = f"{os.path.join(os.path.dirname(os.path.dirname(__file__)), 'demo/kb_output')}/{int(time.time())}.jpg"
        # 如果设置了自定义KB大小
        if idphoto_json["custom_image_kb"]:
            resize_image_to_kb(
                result_image_standard,
                output_image_path,
                idphoto_json["custom_image_kb"],
                dpi=(
                    idphoto_json["custom_image_dpi"]
                    if idphoto_json["custom_image_dpi"]
                    else 300
                ),
            )
            return output_image_path
        # 如果只设置了dpi
        elif idphoto_json["custom_image_dpi"]:
            save_image_dpi_to_bytes(
                result_image_standard,
                output_image_path,
                dpi=idphoto_json["custom_image_dpi"],
            )
            return output_image_path

        return None

    def _create_response(
        self,
        result_image_standard,
        result_image_hd,
        result_image_standard_png,
        result_image_hd_png,
        result_layout_image,
        output_image_path,
    ):
        """创建响应"""
        response = [
            result_image_standard,
            result_image_hd,
            result_image_standard_png,
            result_image_hd_png,
            result_layout_image,
            gr.update(visible=False),
        ]
        if output_image_path:
            response.append(gr.update(visible=True, value=output_image_path))
        else:
            response.append(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,
        ]