import os import gradio as gr from hivision import IDCreator from hivision.error import FaceError 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.human_matting import ( extract_human_modnet_photographic_portrait_matting, extract_human, ) import pathlib import numpy as np from demo.utils import csv_to_size_list import argparse # 获取尺寸列表 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), } color_list_dict_EN = { "Blue": (86, 140, 212), "White": (255, 255, 255), "Red": (233, 51, 35), } # 检测 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, head_measure_ratio=0.2, head_height_ratio=0.45, top_distance_max=0.12, top_distance_min=0.10, ): 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": "人脸数量不等于 1", "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": "The number of faces is not equal to 1", "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() if matting_model_option == "modnet_photographic_portrait_matting": creator.matting_handler = extract_human_modnet_photographic_portrait_matting else: creator.matting_handler = extract_human 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, ) except FaceError: result_message = { img_output_standard: gr.update(value=None), img_output_standard_hd: gr.update(value=None), notification: gr.update( value=text_lang_map[language]["The number of faces is not equal to 1"], 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], ) 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 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 if __name__ == "__main__": # argparser = argparse.ArgumentParser() # argparser.add_argument( # "--port", type=int, default=7860, help="The port number of the server" # ) # argparser.add_argument( # "--host", type=str, default="127.0.0.1", help="The host of the server" # ) # 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") ] 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", "Custom Color"] 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(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", ) matting_model_options = gr.Dropdown( choices=matting_model_list, label="Matting Model", value="hivision_modnet", elem_id="matting_model", ) 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", ) # 左：输出 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, ) 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) img_output_standard_hd = gr.Image(label="高清照", height=350) img_output_layout = gr.Image(label="六寸排版照", height=350) 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="不设置", ), 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 大小后的照片"), } 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", ), 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" ), } 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, custom_image_kb_size, notification, img_output_standard, img_output_standard_hd, img_output_layout, file_download, ], ) 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, ], outputs=[ img_output_standard, img_output_standard_hd, img_output_layout, notification, file_download, ], ) demo.launch()