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 ]