update app.py

app.py

@@ -15,10 +15,17 @@ HUMAN_MATTING_MODELS_EXIST = [
     if file.endswith(".onnx") or file.endswith(".mnn")
 ]
 # 在HUMAN_MATTING_MODELS中的模型才会被加载到Gradio中显示
-HUMAN_MATTING_MODELS = [
+HUMAN_MATTING_MODELS_CHOICE = [
     model for model in HUMAN_MATTING_MODELS if model in HUMAN_MATTING_MODELS_EXIST
 ]
 
+if len(HUMAN_MATTING_MODELS_CHOICE) == 0:
+    raise ValueError(
+        "未找到任何存在的人像分割模型，请检查 hivision/creator/weights 目录下的文件"
+        + "\n"
+        + "No existing portrait segmentation model was found, please check the files in the hivision/creator/weights directory."
+    )
+
 FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn"]
 FACE_DETECT_MODELS_EXPAND = (
     ["retinaface-resnet50"]
@@ -29,7 +36,7 @@ FACE_DETECT_MODELS_EXPAND = (
     )
     else []
 )
-FACE_DETECT_MODELS += FACE_DETECT_MODELS_EXPAND
+FACE_DETECT_MODELS_CHOICE = FACE_DETECT_MODELS + FACE_DETECT_MODELS_EXPAND
 
 LANGUAGE = ["zh", "en", "ko", "ja"]
 
@@ -54,8 +61,8 @@ if __name__ == "__main__":
     demo = create_ui(
         processor,
         root_dir,
-        HUMAN_MATTING_MODELS_EXIST,
-        FACE_DETECT_MODELS,
+        HUMAN_MATTING_MODELS_CHOICE,
+        FACE_DETECT_MODELS_CHOICE,
         LANGUAGE,
     )
     demo.launch(

demo/processor.py

@@ -16,7 +16,7 @@ from demo.utils import range_check
 import gradio as gr
 import os
 import time
-from demo.locals import LOCALES
+from demo.locales import LOCALES
 
 
 class IDPhotoProcessor:
@@ -261,7 +261,7 @@ class IDPhotoProcessor:
         )
 
         # 生成排版照片
-        result_layout_image = self._generate_layout_image(
+        result_image_layout, result_image_layout_gr = self._generate_image_layout(
             idphoto_json,
             result_image_standard,
             language,
@@ -289,7 +289,10 @@ class IDPhotoProcessor:
 
         # 调整图片大小
         output_image_path = self._resize_image_if_needed(
-            result_image_standard, idphoto_json
+            result_image_standard,
+            result_image_hd,
+            result_image_layout,
+            idphoto_json,
         )
 
         return self._create_response(
@@ -297,7 +300,7 @@ class IDPhotoProcessor:
             result_image_hd,
             result_image_standard_png,
             result_image_hd_png,
-            result_layout_image,
+            result_image_layout_gr,
             output_image_path,
         )
 
@@ -319,7 +322,7 @@ class IDPhotoProcessor:
 
         return result_image_standard, result_image_hd
 
-    def _generate_layout_image(
+    def _generate_image_layout(
         self,
         idphoto_json,
         result_image_standard,
@@ -353,14 +356,16 @@ class IDPhotoProcessor:
                 color=watermark_text_color,
             )
 
-        return gr.update(
-            value=generate_layout_image(
-                image,
-                typography_arr,
-                typography_rotate,
-                height=idphoto_json["size"][0],
-                width=idphoto_json["size"][1],
-            ),
+        result_image_layout = generate_layout_image(
+            image,
+            typography_arr,
+            typography_rotate,
+            height=idphoto_json["size"][0],
+            width=idphoto_json["size"][1],
+        )
+
+        return result_image_layout, gr.update(
+            value=result_image_layout,
             visible=True,
         )
 
@@ -390,31 +395,79 @@ class IDPhotoProcessor:
         result_image_hd = add_watermark(image=result_image_hd, **watermark_params)
         return result_image_standard, result_image_hd
 
-    def _resize_image_if_needed(self, result_image_standard, idphoto_json):
+    def _resize_image_if_needed(
+        self,
+        result_image_standard,
+        result_image_hd,
+        result_image_layout,
+        idphoto_json,
+    ):
         """如果需要，调整图片大小"""
-        output_image_path = f"{os.path.join(os.path.dirname(os.path.dirname(__file__)), 'demo/kb_output')}/{int(time.time())}.jpg"
-        # 如果设置了自定义KB大小
-        if idphoto_json["custom_image_kb"]:
+        # 设置输出路径
+        base_path = os.path.join(
+            os.path.dirname(os.path.dirname(__file__)), "demo/kb_output"
+        )
+        timestamp = int(time.time())
+        output_paths = {
+            "standard": f"{base_path}/{timestamp}_standard",
+            "hd": f"{base_path}/{timestamp}_hd",
+            "layout": f"{base_path}/{timestamp}_layout",
+        }
+
+        # 获取自定义的KB和DPI值
+        custom_kb = idphoto_json.get("custom_image_kb")
+        custom_dpi = idphoto_json.get("custom_image_dpi", 300)
+
+        # 处理同时有自定义KB和DPI的情况
+        if custom_kb and custom_dpi:
+            # 为所有输出路径添加DPI信息
+            for key in output_paths:
+                output_paths[key] += f"_{custom_dpi}dpi.jpg"
+            # 为标准图像添加KB信息
+            output_paths["standard"] = output_paths["standard"].replace(
+                ".jpg", f"_{custom_kb}kb.jpg"
+            )
+
+            # 调整标准图像大小并保存
             resize_image_to_kb(
                 result_image_standard,
-                output_image_path,
-                idphoto_json["custom_image_kb"],
-                dpi=(
-                    idphoto_json["custom_image_dpi"]
-                    if idphoto_json["custom_image_dpi"]
-                    else 300
-                ),
+                output_paths["standard"],
+                custom_kb,
+                dpi=custom_dpi,
             )
-            return output_image_path
-        # 如果只设置了dpi
-        elif idphoto_json["custom_image_dpi"]:
+            # 保存高清图像和排版图像
+            save_image_dpi_to_bytes(result_image_hd, output_paths["hd"], dpi=custom_dpi)
             save_image_dpi_to_bytes(
+                result_image_layout, output_paths["layout"], dpi=custom_dpi
+            )
+
+            return list(output_paths.values())
+
+        # 只有自定义DPI的情况
+        elif custom_dpi:
+            for key in output_paths:
+                output_paths[key] += f"_{custom_dpi}dpi.jpg"
+                # 保存所有图像，使用自定义DPI
+                save_image_dpi_to_bytes(
+                    locals()[f"result_image_{key}"], output_paths[key], dpi=custom_dpi
+                )
+
+            return list(output_paths.values())
+
+        # 只有自定义KB的情况
+        elif custom_kb:
+            output_paths["standard"] += f"_{custom_kb}kb.jpg"
+            # 只调整标准图像大小并保存
+            resize_image_to_kb(
                 result_image_standard,
-                output_image_path,
-                dpi=idphoto_json["custom_image_dpi"],
+                output_paths["standard"],
+                custom_kb,
+                dpi=300,
             )
-            return output_image_path
 
+            return [output_paths["standard"]]
+
+        # 如果没有自定义设置，返回None
         return None
 
     def _create_response(

demo/ui.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import pathlib
-from demo.locals import LOCALES
+from demo.locales import LOCALES
 from demo.processor import IDPhotoProcessor
 
 """
@@ -23,7 +23,7 @@ def create_ui(
     face_detect_models: list,
     language: list,
 ):
-    DEFAULT_LANG = language[0]
+    DEFAULT_LANG = "en"
     DEFAULT_HUMAN_MATTING_MODEL = "modnet_photographic_portrait_matting"
     DEFAULT_FACE_DETECT_MODEL = "retinaface-resnet50"
 

demo/utils.py

@@ -1,7 +1,4 @@
 import csv
-import numpy as np
-from PIL import Image
-from hivision.plugin.watermark import Watermarker, WatermarkerStyles
 
 
 def csv_to_size_list(csv_file: str) -> dict:

hivision/creator/human_matting.py

@@ -37,10 +37,9 @@ WEIGHTS = {
     ),
 }
 
-ONNX_DEVICE = (
-    "CUDAExecutionProvider"
-    if onnxruntime.get_device() == "GPU"
-    else "CPUExecutionProvider"
+ONNX_DEVICE = onnxruntime.get_device()
+ONNX_PROVIDER = (
+    "CUDAExecutionProvider" if ONNX_DEVICE == "GPU" else "CPUExecutionProvider"
 )
 
 HIVISION_MODNET_SESS = None
@@ -52,7 +51,7 @@ BIREFNET_V1_LITE_SESS = None
 def load_onnx_model(checkpoint_path, set_cpu=False):
     providers = (
         ["CUDAExecutionProvider", "CPUExecutionProvider"]
-        if ONNX_DEVICE == "CUDAExecutionProvider"
+        if ONNX_PROVIDER == "CUDAExecutionProvider"
         else ["CPUExecutionProvider"]
     )
 
@@ -365,7 +364,17 @@ def get_birefnet_portrait_matting(input_image, checkpoint_path, ref_size=512):
 
     if BIREFNET_V1_LITE_SESS is None:
         print("首次加载birefnet-v1-lite模型...")
-        BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path)
+        if ONNX_DEVICE == "GPU":
+            print("onnxruntime-gpu已安装，尝试使用CUDA加载模型")
+            try:
+                import torch
+            except ImportError:
+                print(
+                    "torch未安装，尝试直接使用onnxruntime-gpu加载模型，这需要配置好CUDA和cuDNN"
+                )
+            BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path)
+        else:
+            BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path, set_cpu=True)
 
     # 记录加载onnx模型的结束时间
     load_end_time = time()