diff --git a/README.md b/README.md index 1b126ac..4fd3301 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ - 在线体验： [![SwanHub Demo](https://img.shields.io/static/v1?label=Demo&message=SwanHub%20Demo&color=blue)](https://swanhub.co/ZeYiLin/HivisionIDPhotos/demo)、[![Spaces](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue)](https://huggingface.co/spaces/TheEeeeLin/HivisionIDPhotos) +- 2024.9.6: 增加新的抠图模型 [modnet_photographic_portrait_matting.onnx](https://github.com/ZHKKKe/MODNet) - 2024.9.5: 更新 [Restful API 文档](docs/api_CN.md) - 2024.9.2: 更新**调整照片 KB 大小**，[DockerHub](https://hub.docker.com/r/linzeyi/hivision_idphotos/tags) - 2023.12.1: 更新**API 部署（基于 fastapi）** @@ -88,6 +89,10 @@ pip install -r requirements-app.txt 在我们的[Release](https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/tag/pretrained-model)下载权重文件`hivision_modnet.onnx` (24.7MB)，存到项目的`hivision/creator/weights`目录下。 +拓展抠图模型权重（均放到`hivision/creator/weights`目录下）： +- modnet_photographic_portrait_matting.onnx: [MODNet](https://github.com/ZHKKKe/MODNet)官方权重，[下载](https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR) + +
# 🚀 运行 Gradio Demo diff --git a/README_EN.md b/README_EN.md index b9cac80..15c3ccf 100644 --- a/README_EN.md +++ b/README_EN.md @@ -21,6 +21,7 @@ English / [中文](README.md) / [日本語](README_JP.md) / [한국어](README_K - Online Demo: [![SwanHub Demo](https://img.shields.io/static/v1?label=Demo&message=SwanHub%20Demo&color=blue)](https://swanhub.co/ZeYiLin/HivisionIDPhotos/demo)、[![Spaces](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue)](https://huggingface.co/spaces/TheEeeeLin/HivisionIDPhotos) +- 2024.9.6: Add a new image matting model [modnet_photographic_portrait_matting.onnx](https://github.com/ZHKKKe/MODNet) - 2024.9.2: Update **Adjusted photo KB size**，[DockerHub](https://hub.docker.com/r/linzeyi/hivision_idphotos/tags) - 2023.12.1: Update **API deployment (based on fastapi)** - 2023.6.20: Update **Preset size menu** @@ -80,6 +81,10 @@ pip install -r requirements-app.txt Download the weight file `hivision_modnet.onnx` from our [Release](https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/tag/pretrained-model) and save it to the `hivision/creator/weights` directory. +Expand matting model weights (all in the `hivision/creator/weights` directory) : + +- modnet_photographic_portrait_matting.onnx: by [MODNet](https://github.com/ZHKKKe/MODNet)，[Download](https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR) +
# 🚀 Gradio Demo diff --git a/app.py b/app.py index 9ce7f57..1e7950c 100644 --- a/app.py +++ b/app.py @@ -7,6 +7,10 @@ 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 @@ -54,6 +58,7 @@ def idphoto_inference( 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, @@ -146,6 +151,11 @@ def idphoto_inference( 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: @@ -267,7 +277,24 @@ def idphoto_inference( 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"] @@ -283,14 +310,6 @@ if __name__ == "__main__": image_kb_CN = ["不设置", "自定义"] image_kb_EN = ["Not Set", "Custom"] - # title = "

HivisionIDPhotos

" - # description = "

😎9.2 Update: Add photo size KB adjustment

" - # css = """ - # h1#title, h3 { - # text-align: center; - # } - # """ - css = """ #col-left { margin: 0 auto; @@ -326,9 +345,20 @@ if __name__ == "__main__": # ------------ 左半边 UI ---------------- with gr.Column(): img_input = gr.Image(height=400) - language_options = gr.Dropdown( - choices=language, label="Language", value="中文", elem_id="language" - ) + + 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, @@ -453,6 +483,7 @@ if __name__ == "__main__": img_output_layout: gr.update(label="六寸排版照"), file_download: gr.update(label="下载调整 KB 大小后的照片"), } + elif language == "English": return { size_list_options: gr.update( @@ -576,6 +607,7 @@ if __name__ == "__main__": custom_size_wdith, custom_image_kb_size, language_options, + matting_model_options, ], outputs=[ img_output_standard, @@ -586,13 +618,4 @@ 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() - demo.launch(server_name=args.host, server_port=args.port) diff --git a/deploy_api.py b/deploy_api.py index 61a77ad..7b03d28 100644 --- a/deploy_api.py +++ b/deploy_api.py @@ -146,12 +146,5 @@ async def generate_layout_photos( if __name__ == "__main__": import uvicorn - # 加载权重文件 - root_dir = os.path.dirname(os.path.abspath(__file__)) - HY_HUMAN_MATTING_WEIGHTS_PATH = os.path.join( - root_dir, "hivision/creator/weights/hivision_modnet.onnx" - ) - sess = onnxruntime.InferenceSession(HY_HUMAN_MATTING_WEIGHTS_PATH) - # 在8080端口运行推理服务 uvicorn.run(app, host="0.0.0.0", port=8080) diff --git a/hivision/creator/human_matting.py b/hivision/creator/human_matting.py index 4eac05c..dfa71f9 100644 --- a/hivision/creator/human_matting.py +++ b/hivision/creator/human_matting.py @@ -15,7 +15,17 @@ from .context import Context import cv2 import os -weight_path = os.path.join(os.path.dirname(__file__), "weights", "hivision_modnet.onnx") + +WEIGHTS = { + "hivision_modnet": os.path.join( + os.path.dirname(__file__), "weights", "hivision_modnet.onnx" + ), + "modnet_photographic_portrait_matting": os.path.join( + os.path.dirname(__file__), + "weights", + "modnet_photographic_portrait_matting.onnx", + ), +} def extract_human(ctx: Context): @@ -24,7 +34,21 @@ def extract_human(ctx: Context): :param ctx: 上下文 """ # 抠图 - matting_image = get_modnet_matting(ctx.processing_image, weight_path) + matting_image = get_modnet_matting(ctx.processing_image, WEIGHTS["hivision_modnet"]) + # 修复抠图 + ctx.processing_image = hollow_out_fix(matting_image) + ctx.matting_image = ctx.processing_image.copy() + + +def extract_human_modnet_photographic_portrait_matting(ctx: Context): + """ + 人像抠图 + :param ctx: 上下文 + """ + # 抠图 + matting_image = get_modnet_matting( + ctx.processing_image, WEIGHTS["modnet_photographic_portrait_matting"] + ) # 修复抠图 ctx.processing_image = hollow_out_fix(matting_image) ctx.matting_image = ctx.processing_image.copy() @@ -92,13 +116,13 @@ def read_modnet_image(input_image, ref_size=512): return im, width, length -sess = None +# sess = None def get_modnet_matting(input_image, checkpoint_path, ref_size=512): - global sess - if sess is None: - sess = onnxruntime.InferenceSession(checkpoint_path) + # global sess + # if sess is None: + sess = onnxruntime.InferenceSession(checkpoint_path) input_name = sess.get_inputs()[0].name output_name = sess.get_outputs()[0].name diff --git a/inference.py b/inference.py index 4096334..53c7196 100644 --- a/inference.py +++ b/inference.py @@ -34,13 +34,6 @@ args = parser.parse_args() root_dir = os.path.dirname(os.path.abspath(__file__)) -# 预加载 ONNX 模型 -print("正在加载抠图模型...") -# HY_HUMAN_MATTING_WEIGHTS_PATH = os.path.join( -# root_dir, "hivision/creator/weights/hivision_modnet.onnx" -# ) -# sess = onnxruntime.InferenceSession(HY_HUMAN_MATTING_WEIGHTS_PATH) - input_image = cv2.imread(args.input_image_dir, cv2.IMREAD_UNCHANGED)