update new model

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 assets/demoImage.png filter=lfs diff=lfs merge=lfs -text
+hivision/creator/weights/rmbg-1.4.onnx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,10 +1,11 @@
 *.pyc
-**/__pycache__
+**/__pycache__/
 .idea
 .vscode/*
 .DS_Store
-app/output/*.jpg
+.env
 demo/kb_output/*.jpg
+demo/kb_output/*.png
 # build outputs
 dist
 build
@@ -12,5 +13,9 @@ build
 *.pth
 *.pt
 *.onnx
+*.mnn
 test/temp/*
-!test/temp/.gitkeep
+!test/temp/.gitkeep
+!hivision/creator/weights/rmbg-1.4.onnx
+
+.python-version

app.py

@@ -7,8 +7,7 @@ from hivision.creator.layout_calculator import (
     generate_layout_photo,
     generate_layout_image,
 )
-from hivision.creator.human_matting import *
-from hivision.creator.face_detector import *
+from hivision.creator.choose_handler import choose_handler
 import pathlib
 import numpy as np
 from demo.utils import csv_to_size_list
@@ -150,15 +149,7 @@ 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
-
-    if face_detect_option == "mtcnn":
-        creator.detection_handler = detect_face_mtcnn
-    else:
-        creator.detection_handler = detect_face_face_plusplus
+    choose_handler(creator, matting_model_option, face_detect_option)
 
     change_bg_only = idphoto_json["size_mode"] in ["只换底", "Only Change Background"]
     # 生成证件照
@@ -294,28 +285,28 @@ 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"
-    # )
-    # argparser.add_argument(
-    #     "--root_path",
-    #     type=str,
-    #     default=None,
-    #     help="The root path of the server, default is None (='/'), e.g. '/myapp'",
-    # )
-
-    # args = argparser.parse_args()
+    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"
+    )
+    argparser.add_argument(
+        "--root_path",
+        type=str,
+        default=None,
+        help="The root path of the server, default is None (='/'), e.g. '/myapp'",
+    )
+
+    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")
+        if file.endswith(".onnx") or file.endswith(".mnn")
     ]
     DEFAULT_MATTING_MODEL = "modnet_photographic_portrait_matting"
     if DEFAULT_MATTING_MODEL in matting_model_list:
@@ -366,7 +357,7 @@ if __name__ == "__main__":
             content = f.read()
         return content
 
-    demo = gr.Blocks(css=css)
+    demo = gr.Blocks(title="HivisionIDPhotos", css=css)
 
     with demo:
         gr.HTML(load_description(os.path.join(root_dir, "assets/title.md")))
@@ -669,7 +660,7 @@ if __name__ == "__main__":
     demo.launch(
         # server_name=args.host,
         # server_port=args.port,
-        # show_api=False,
+        show_api=False,
         # favicon_path=os.path.join(root_dir, "assets/hivision_logo.png"),
         # root_path=args.root_path,
     )

hivision/creator/__init__.py

@@ -41,8 +41,7 @@ class IDCreator:
         """
         # 处理者
         self.matting_handler: ContextHandler = extract_human
-        self.detection_handler: ContextHandler = detect_face_face_plusplus
-        # self.detection_handler: ContextHandler = detect_face_mtcnn
+        self.detection_handler: ContextHandler = detect_face_mtcnn
 
         # 上下文
         self.ctx = None

hivision/creator/choose_handler.py

@@ -0,0 +1,18 @@
+from hivision.creator.human_matting import *
+from hivision.creator.face_detector import *
+
+
+def choose_handler(creator, matting_model_option=None, face_detect_option=None):
+    if matting_model_option == "modnet_photographic_portrait_matting":
+        creator.matting_handler = extract_human_modnet_photographic_portrait_matting
+    elif matting_model_option == "mnn_hivision_modnet":
+        creator.matting_handler = extract_human_mnn_modnet
+    elif matting_model_option == "rmbg-1.4":
+        creator.matting_handler = extract_human_rmbg
+    else:
+        creator.matting_handler = extract_human
+
+    if face_detect_option == "face_plusplus":
+        creator.detection_handler = detect_face_face_plusplus
+    else:
+        creator.detection_handler = detect_face_mtcnn

hivision/creator/human_matting.py

@@ -25,8 +25,43 @@ WEIGHTS = {
         "weights",
         "modnet_photographic_portrait_matting.onnx",
     ),
+    "mnn_hivision_modnet": os.path.join(
+        os.path.dirname(__file__),
+        "weights",
+        "mnn_hivision_modnet.mnn",
+    ),
+    "rmbg-1.4": os.path.join(os.path.dirname(__file__), "weights", "rmbg-1.4.onnx"),
 }
 
+ONNX_DEVICE = (
+    "CUDAExecutionProvider"
+    if onnxruntime.get_device() == "GPU"
+    else "CPUExecutionProvider"
+)
+
+
+def load_onnx_model(checkpoint_path):
+    providers = (
+        ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        if ONNX_DEVICE == "CUDAExecutionProvider"
+        else ["CPUExecutionProvider"]
+    )
+
+    try:
+        sess = onnxruntime.InferenceSession(checkpoint_path, providers=providers)
+    except Exception as e:
+        if ONNX_DEVICE == "CUDAExecutionProvider":
+            print(f"Failed to load model with CUDAExecutionProvider: {e}")
+            print("Falling back to CPUExecutionProvider")
+            # 尝试使用CPU加载模型
+            sess = onnxruntime.InferenceSession(
+                checkpoint_path, providers=["CPUExecutionProvider"]
+            )
+        else:
+            raise e  # 如果是CPU执行失败，重新抛出异常
+
+    return sess
+
 
 def extract_human(ctx: Context):
     """
@@ -50,10 +85,24 @@ def extract_human_modnet_photographic_portrait_matting(ctx: Context):
         ctx.processing_image, WEIGHTS["modnet_photographic_portrait_matting"]
     )
     # 修复抠图
+    ctx.processing_image = matting_image
+    ctx.matting_image = ctx.processing_image.copy()
+
+
+def extract_human_mnn_modnet(ctx: Context):
+    matting_image = get_mnn_modnet_matting(
+        ctx.processing_image, WEIGHTS["mnn_hivision_modnet"]
+    )
     ctx.processing_image = hollow_out_fix(matting_image)
     ctx.matting_image = ctx.processing_image.copy()
 
 
+def extract_human_rmbg(ctx: Context):
+    matting_image = get_rmbg_matting(ctx.processing_image, WEIGHTS["rmbg-1.4"])
+    ctx.processing_image = matting_image
+    ctx.matting_image = ctx.processing_image.copy()
+
+
 def hollow_out_fix(src: np.ndarray) -> np.ndarray:
     """
     修补抠图区域，作为抠图模型精度不够的补充
@@ -120,9 +169,11 @@ def read_modnet_image(input_image, ref_size=512):
 
 
 def get_modnet_matting(input_image, checkpoint_path, ref_size=512):
-    # global sess
-    # if sess is None:
-    sess = onnxruntime.InferenceSession(checkpoint_path)
+    if not os.path.exists(checkpoint_path):
+        print(f"Checkpoint file not found: {checkpoint_path}")
+        return None
+
+    sess = load_onnx_model(checkpoint_path)
 
     input_name = sess.get_inputs()[0].name
     output_name = sess.get_outputs()[0].name
@@ -138,3 +189,85 @@ def get_modnet_matting(input_image, checkpoint_path, ref_size=512):
     output_image = cv2.merge((b, g, r, mask))
 
     return output_image
+
+
+def get_rmbg_matting(input_image: np.ndarray, checkpoint_path, ref_size=1024):
+    if not os.path.exists(checkpoint_path):
+        print(f"Checkpoint file not found: {checkpoint_path}")
+        return None
+
+    def resize_rmbg_image(image):
+        image = image.convert("RGB")
+        model_input_size = (ref_size, ref_size)
+        image = image.resize(model_input_size, Image.BILINEAR)
+        return image
+
+    sess = load_onnx_model(checkpoint_path)
+
+    orig_image = Image.fromarray(input_image)
+    image = resize_rmbg_image(orig_image)
+    im_np = np.array(image).astype(np.float32)
+    im_np = im_np.transpose(2, 0, 1)  # Change to CxHxW format
+    im_np = np.expand_dims(im_np, axis=0)  # Add batch dimension
+    im_np = im_np / 255.0  # Normalize to [0, 1]
+    im_np = (im_np - 0.5) / 0.5  # Normalize to [-1, 1]
+
+    # Inference
+    result = sess.run(None, {sess.get_inputs()[0].name: im_np})[0]
+
+    # Post process
+    result = np.squeeze(result)
+    ma = np.max(result)
+    mi = np.min(result)
+    result = (result - mi) / (ma - mi)  # Normalize to [0, 1]
+
+    # Convert to PIL image
+    im_array = (result * 255).astype(np.uint8)
+    pil_im = Image.fromarray(
+        im_array, mode="L"
+    )  # Ensure mask is single channel (L mode)
+
+    # Resize the mask to match the original image size
+    pil_im = pil_im.resize(orig_image.size, Image.BILINEAR)
+
+    # Paste the mask on the original image
+    new_im = Image.new("RGBA", orig_image.size, (0, 0, 0, 0))
+    new_im.paste(orig_image, mask=pil_im)
+
+    return np.array(new_im)
+
+
+def get_mnn_modnet_matting(input_image, checkpoint_path, ref_size=512):
+    if not os.path.exists(checkpoint_path):
+        print(f"Checkpoint file not found: {checkpoint_path}")
+        return None
+
+    try:
+        import MNN.expr as expr
+        import MNN.nn as nn
+    except ImportError as e:
+        raise ImportError(
+            "The MNN module is not installed or there was an import error. Please ensure that the MNN library is installed by using the command 'pip install mnn'."
+        ) from e
+
+    config = {}
+    config["precision"] = "low"  # 当硬件支持（armv8.2）时使用fp16推理
+    config["backend"] = 0  # CPU
+    config["numThread"] = 4  # 线程数
+    im, width, length = read_modnet_image(input_image, ref_size=512)
+    rt = nn.create_runtime_manager((config,))
+    net = nn.load_module_from_file(
+        checkpoint_path, ["input1"], ["output1"], runtime_manager=rt
+    )
+    input_var = expr.convert(im, expr.NCHW)
+    output_var = net.forward(input_var)
+    matte = expr.convert(output_var, expr.NCHW)
+    matte = matte.read()  # var转换为np
+    matte = (matte * 255).astype("uint8")
+    matte = np.squeeze(matte)
+    mask = cv2.resize(matte, (width, length), interpolation=cv2.INTER_AREA)
+    b, g, r = cv2.split(np.uint8(input_image))
+
+    output_image = cv2.merge((b, g, r, mask))
+
+    return output_image

hivision/creator/weights/rmbg-1.4.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cafcf770b06757c4eaced21b1a88e57fd2b66de01b8045f35f01535ba742e0f
+size 176153355