diff --git a/deploy_api.py b/deploy_api.py index 7b03d28..92b9cbf 100644 --- a/deploy_api.py +++ b/deploy_api.py @@ -10,7 +10,6 @@ from hivision.utils import add_background, resize_image_to_kb_base64, hex_to_rgb import base64 import numpy as np import cv2 -import os app = FastAPI() creator = IDCreator() @@ -65,8 +64,13 @@ async def idphoto_inference( # 透明图像添加纯色背景接口 @app.post("/add_background") async def photo_add_background( - input_image: UploadFile, color: str = Form(...), kb: str = Form(None) + input_image: UploadFile, + color: str = Form(...), + kb: str = Form(None), + render: int = Form(0), ): + render_choice = ["pure_color", "updown_gradient", "center_gradient"] + image_bytes = await input_image.read() nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED) @@ -74,7 +78,11 @@ async def photo_add_background( color = hex_to_rgb(color) color = (color[2], color[1], color[0]) - result_image = add_background(img, bgr=color).astype(np.uint8) + result_image = add_background( + img, + bgr=color, + mode=render_choice[render], + ).astype(np.uint8) if kb: result_image = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR) @@ -83,7 +91,6 @@ async def photo_add_background( result_image_base64 = numpy_2_base64(result_image) # try: - result_messgae = { "status": True, "image_base64": result_image_base64, diff --git a/inference.py b/inference.py index 5656928..e2a04e4 100644 --- a/inference.py +++ b/inference.py @@ -25,7 +25,11 @@ INFERENCE_TYPE = [ "add_background", "generate_layout_photos", ] -MATTING_MODEL = ["hivision_modnet", "modnet_photographic_portrait_matting", "mnn_hivision_modnet"] +MATTING_MODEL = [ + "hivision_modnet", + "modnet_photographic_portrait_matting", + "mnn_hivision_modnet", +] RENDER = [0, 1, 2] parser.add_argument( diff --git a/requests_api.py b/requests_api.py index ac3f8de..bf35b15 100644 --- a/requests_api.py +++ b/requests_api.py @@ -3,6 +3,19 @@ import base64 import argparse import os +INFERENCE_TYPE = [ + "idphoto", + "human_matting", + "add_background", + "generate_layout_photos", +] +MATTING_MODEL = [ + "hivision_modnet", + "modnet_photographic_portrait_matting", + "mnn_hivision_modnet", +] +RENDER = [0, 1, 2] + def base64_save(_base64_image_data, save_path): # 解码 Base64 数据并保存为 PNG 文件 @@ -66,16 +79,21 @@ if __name__ == "__main__": parser.add_argument( "-k", "--kb", help="输出照片的 KB 值,仅对换底和制作排版照生效", default=None ) + parser.add_argument( + "-r", + "--render", + type=int, + help="底色合成的模式,有 0:纯色、1:上下渐变、2:中心渐变 可选", + choices=RENDER, + default=0, + ) args = parser.parse_args() - url = f"{args.url}/{args.type}" # 替换为实际的接口 URL - # color = hex_to_rgb(args.color) - # color = (color[2], color[1], color[0]) if args.type == "idphoto": # 调用 /idphoto 接口 idphoto_response = request_idphoto( - args.input_image_dir, int(args.height), int(args.width) + args.input_image_dir, int(args.height), int(args.width), args.render ) if idphoto_response["status"]: