fix:bg-bug

.gitignore

@@ -3,8 +3,7 @@
 .idea
 .vscode/*
 .DS_Store
-hivision_modnet.onnx
-output/*.jpg
+app/output/*.jpg
 # build outputs
 dist
 build

app/web.py

@@ -1,16 +1,17 @@
 import os
 import gradio as gr
-import onnxruntime
-from hivision.creator.face_judgement_align import IDphotos_create
-from hivision.creator.vision import add_background
-from hivision.creator.layoutCreate import generate_layout_photo, generate_layout_image
+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,
+)
 import pathlib
 import numpy as np
-from utils.image_utils import resize_image_to_kb
 from app.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, "size_list_CN.csv"))
@@ -149,45 +150,28 @@ def idphoto_inference(
     else:
         idphoto_json["custom_image_kb"] = None
 
+    creator = IDCreator()
+    change_bg_only = idphoto_json["size_mode"] in ["只换底", "Only Change Background"]
     # 生成证件照
-    (
-        result_image_hd,
-        result_image_standard,
-        typography_arr,
-        typography_rotate,
-        _,
-        _,
-        _,
-        _,
-        status,
-    ) = IDphotos_create(
-        input_image,
-        mode=idphoto_json["size_mode"],
-        size=idphoto_json["size"],
-        head_measure_ratio=head_measure_ratio,
-        head_height_ratio=head_height_ratio,
-        align=False,
-        beauty=False,
-        fd68=None,
-        human_sess=sess,
-        IS_DEBUG=False,
-        top_distance_max=top_distance_max,
-        top_distance_min=top_distance_min,
-    )
-
-    # 如果检测到人脸数量不等于 1
-    if status == 0:
-        result_messgae = {
+    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,
             ),
         }
-
-    # 如果检测到人脸数量等于 1
     else:
+        (result_image_hd, result_image_standard, _, _) = 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"])
@@ -255,7 +239,7 @@ def idphoto_inference(
             # 输出路径为一个根据时间戳 + 哈希值生成的随机文件名
             import time
 
-            output_image_path = f"./output/{int(time.time())}.jpg"
+            output_image_path = f"{os.path.join(os.path.dirname(__file__), 'output')}/{int(time.time())}.jpg"
             resize_image_to_kb(
                 result_image_standard,
                 output_image_path,
@@ -265,30 +249,26 @@ def idphoto_inference(
             output_image_path = None
 
         if output_image_path:
-            result_messgae = {
+            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_messgae = {
+            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_messgae
+    return result_message
 
 
 if __name__ == "__main__":
-    # 预加载 ONNX 模型
-    HY_HUMAN_MATTING_WEIGHTS_PATH = os.path.join(root_dir, "hivision_modnet.onnx")
-    sess = onnxruntime.InferenceSession(HY_HUMAN_MATTING_WEIGHTS_PATH)
-
     language = ["中文", "English"]
     size_mode_CN = ["尺寸列表", "只换底", "自定义尺寸"]
     size_mode_EN = ["Size List", "Only Change Background", "Custom Size"]

hivision/creator/__init__.py

@@ -89,6 +89,15 @@ def __call__(
         # 1. 人像抠图
         self.matting_handler(ctx)
         self.after_matting and self.after_matting(ctx)
+        if ctx.params.change_bg_only:
+            ctx.result = Result(
+                standard=ctx.matting_image,
+                hd=ctx.matting_image,
+                clothing_params=None,
+                typography_params=None,
+            )
+            self.after_all and self.after_all(ctx)
+            return ctx.result
         # 2. 人脸检测
         self.detection_handler(ctx)
         self.after_detect and self.after_detect(ctx)

hivision/creator/context.py

@@ -52,13 +52,24 @@ def __init__(
         self,
         standard: np.ndarray,
         hd: np.ndarray,
-        clothing_params: dict,
-        typography_params: dict,
+        clothing_params: Optional[dict],
+        typography_params: Optional[dict],
     ):
         self.standard = standard
         self.hd = hd
         self.clothing_params = clothing_params
+        """
+        服装参数，仅换底时为 None
+        """
         self.typography_params = typography_params
+        """
+        排版参数，仅换底时为 None
+        """
+
+    def __iter__(self):
+        return iter(
+            [self.standard, self.hd, self.clothing_params, self.typography_params]
+        )
 
 
 class Context:

hivision/creator/human_matting.py

@@ -92,8 +92,13 @@ def read_modnet_image(input_image, ref_size=512):
     return im, width, length
 
 
+sess = None
+
+
 def get_modnet_matting(input_image, checkpoint_path, ref_size=512):
-    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

hivision/utils.py

@@ -168,3 +168,86 @@ def hex_to_rgb(value):
     return tuple(
         int(value[i : i + length // 3], 16) for i in range(0, length, length // 3)
     )
+
+
+def generate_gradient(start_color, width, height, mode="updown"):
+    # 定义背景颜色
+    end_color = (255, 255, 255)  # 白色
+
+    # 创建一个空白图像
+    r_out = np.zeros((height, width), dtype=int)
+    g_out = np.zeros((height, width), dtype=int)
+    b_out = np.zeros((height, width), dtype=int)
+
+    if mode == "updown":
+        # 生成上下渐变色
+        for y in range(height):
+            r = int(
+                (y / height) * end_color[0] + ((height - y) / height) * start_color[0]
+            )
+            g = int(
+                (y / height) * end_color[1] + ((height - y) / height) * start_color[1]
+            )
+            b = int(
+                (y / height) * end_color[2] + ((height - y) / height) * start_color[2]
+            )
+            r_out[y, :] = r
+            g_out[y, :] = g
+            b_out[y, :] = b
+
+    else:
+        # 生成中心渐变色
+        img = np.zeros((height, width, 3))
+        # 定义椭圆中心和半径
+        center = (width // 2, height // 2)
+        end_axies = max(height, width)
+        # 定义渐变色
+        end_color = (255, 255, 255)
+        # 绘制椭圆
+        for y in range(end_axies):
+            axes = (end_axies - y, end_axies - y)
+            r = int(
+                (y / end_axies) * end_color[0]
+                + ((end_axies - y) / end_axies) * start_color[0]
+            )
+            g = int(
+                (y / end_axies) * end_color[1]
+                + ((end_axies - y) / end_axies) * start_color[1]
+            )
+            b = int(
+                (y / end_axies) * end_color[2]
+                + ((end_axies - y) / end_axies) * start_color[2]
+            )
+
+            cv2.ellipse(img, center, axes, 0, 0, 360, (b, g, r), -1)
+        b_out, g_out, r_out = cv2.split(np.uint64(img))
+
+    return r_out, g_out, b_out
+
+
+def add_background(input_image, bgr=(0, 0, 0), mode="pure_color"):
+    """
+    本函数的功能为为透明图像加上背景。
+    :param input_image: numpy.array(4 channels), 透明图像
+    :param bgr: tuple, 合成纯色底时的 BGR 值
+    :param new_background: numpy.array(3 channels)，合成自定义图像底时的背景图
+    :return: output: 合成好的输出图像
+    """
+    height, width = input_image.shape[0], input_image.shape[1]
+    b, g, r, a = cv2.split(input_image)
+    a_cal = a / 255
+    if mode == "pure_color":
+        # 纯色填充
+        b2 = np.full([height, width], bgr[0], dtype=int)
+        g2 = np.full([height, width], bgr[1], dtype=int)
+        r2 = np.full([height, width], bgr[2], dtype=int)
+    elif mode == "updown_gradient":
+        b2, g2, r2 = generate_gradient(bgr, width, height, mode="updown")
+    else:
+        b2, g2, r2 = generate_gradient(bgr, width, height, mode="center")
+
+    output = cv2.merge(
+        ((b - b2) * a_cal + b2, (g - g2) * a_cal + g2, (r - r2) * a_cal + r2)
+    )
+
+    return output

requirements-app.txt

@@ -1,2 +1,2 @@
-gradio>=4.37.1
+gradio==3.38.0
 fastapi