update

.gitattributes

@@ -1,37 +1 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip 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
+hivision/creator/weights/birefnet-v1-lite.onnx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -17,5 +17,6 @@ build
 test/temp/*
 !test/temp/.gitkeep
 !hivision/creator/weights/rmbg-1.4.onnx
+!hivision/creator/weights/birefnet-v1-lite.onnx
 
 .python-version

app.py

@@ -444,7 +444,7 @@ if __name__ == "__main__":
                         minimum=0.1,
                         maximum=0.5,
                         value=0.2,
-                        step=0.02,
+                        step=0.01,
                         label="面部比例",
                         interactive=True,
                     )
@@ -453,7 +453,7 @@ if __name__ == "__main__":
                         minimum=0.02,
                         maximum=0.5,
                         value=0.12,
-                        step=0.02,
+                        step=0.01,
                         label="头距顶距离",
                         interactive=True,
                     )

hivision/.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip 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

hivision/creator/choose_handler.py

@@ -9,10 +9,17 @@ def choose_handler(creator, matting_model_option=None, face_detect_option=None):
         creator.matting_handler = extract_human_mnn_modnet
     elif matting_model_option == "rmbg-1.4":
         creator.matting_handler = extract_human_rmbg
+    # elif matting_model_option == "birefnet-portrait":
+    #     creator.matting_handler = extract_human_birefnet_portrait
+    elif matting_model_option == "birefnet-v1-lite":
+        creator.matting_handler = extract_human_birefnet_lite
     else:
         creator.matting_handler = extract_human
 
-    if face_detect_option == "face_plusplus":
+    if (
+        face_detect_option == "face_plusplus"
+        or face_detect_option == "face++ (联网API)"
+    ):
         creator.detection_handler = detect_face_face_plusplus
     else:
         creator.detection_handler = detect_face_mtcnn

hivision/creator/face_detector.py

@@ -65,6 +65,8 @@ def detect_face_face_plusplus(ctx: Context):
     api_key = os.getenv("FACE_PLUS_API_KEY")
     api_secret = os.getenv("FACE_PLUS_API_SECRET")
 
+    print("调用了face++")
+
     image = ctx.origin_image
     # 将图片转为 base64, 且不大于2MB（Face++ API接口限制）
     image_base64 = resize_image_to_kb_base64(image, 2000, mode="max")

hivision/creator/human_matting.py

@@ -14,6 +14,7 @@ from .tensor2numpy import NNormalize, NTo_Tensor, NUnsqueeze
 from .context import Context
 import cv2
 import os
+from time import time
 
 
 WEIGHTS = {
@@ -31,6 +32,9 @@ WEIGHTS = {
         "mnn_hivision_modnet.mnn",
     ),
     "rmbg-1.4": os.path.join(os.path.dirname(__file__), "weights", "rmbg-1.4.onnx"),
+    "birefnet-v1-lite": os.path.join(
+        os.path.dirname(__file__), "weights", "birefnet-v1-lite.onnx"
+    ),
 }
 
 ONNX_DEVICE = (
@@ -39,26 +43,36 @@ ONNX_DEVICE = (
     else "CPUExecutionProvider"
 )
 
+HIVISION_MODNET_SESS = None
+MODNET_PHOTOGRAPHIC_PORTRAIT_MATTING_SESS = None
+RMBG_SESS = None
+BIREFNET_V1_LITE_SESS = None
+
 
-def load_onnx_model(checkpoint_path):
+def load_onnx_model(checkpoint_path, set_cpu=False):
     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执行失败，重新抛出异常
+    if set_cpu:
+        sess = onnxruntime.InferenceSession(
+            checkpoint_path, providers=["CPUExecutionProvider"]
+        )
+    else:
+        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
 
@@ -103,6 +117,22 @@ def extract_human_rmbg(ctx: Context):
     ctx.matting_image = ctx.processing_image.copy()
 
 
+# def extract_human_birefnet_portrait(ctx: Context):
+#     matting_image = get_birefnet_portrait_matting(
+#         ctx.processing_image, WEIGHTS["birefnet-portrait"]
+#     )
+#     ctx.processing_image = matting_image
+#     ctx.matting_image = ctx.processing_image.copy()
+
+
+def extract_human_birefnet_lite(ctx: Context):
+    matting_image = get_birefnet_portrait_matting(
+        ctx.processing_image, WEIGHTS["birefnet-v1-lite"]
+    )
+    ctx.processing_image = matting_image
+    ctx.matting_image = ctx.processing_image.copy()
+
+
 def hollow_out_fix(src: np.ndarray) -> np.ndarray:
     """
     修补抠图区域，作为抠图模型精度不够的补充
@@ -165,22 +195,22 @@ 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):
+    global HIVISION_MODNET_SESS
+
     if not os.path.exists(checkpoint_path):
         print(f"Checkpoint file not found: {checkpoint_path}")
         return None
 
-    sess = load_onnx_model(checkpoint_path)
+    if HIVISION_MODNET_SESS is None:
+        HIVISION_MODNET_SESS = load_onnx_model(checkpoint_path, set_cpu=True)
 
-    input_name = sess.get_inputs()[0].name
-    output_name = sess.get_outputs()[0].name
+    input_name = HIVISION_MODNET_SESS.get_inputs()[0].name
+    output_name = HIVISION_MODNET_SESS.get_outputs()[0].name
 
     im, width, length = read_modnet_image(input_image=input_image, ref_size=ref_size)
 
-    matte = sess.run([output_name], {input_name: im})
+    matte = HIVISION_MODNET_SESS.run([output_name], {input_name: im})
     matte = (matte[0] * 255).astype("uint8")
     matte = np.squeeze(matte)
     mask = cv2.resize(matte, (width, length), interpolation=cv2.INTER_AREA)
@@ -192,6 +222,8 @@ def get_modnet_matting(input_image, checkpoint_path, ref_size=512):
 
 
 def get_rmbg_matting(input_image: np.ndarray, checkpoint_path, ref_size=1024):
+    global RMBG_SESS
+
     if not os.path.exists(checkpoint_path):
         print(f"Checkpoint file not found: {checkpoint_path}")
         return None
@@ -202,7 +234,8 @@ def get_rmbg_matting(input_image: np.ndarray, checkpoint_path, ref_size=1024):
         image = image.resize(model_input_size, Image.BILINEAR)
         return image
 
-    sess = load_onnx_model(checkpoint_path)
+    if RMBG_SESS is None:
+        RMBG_SESS = load_onnx_model(checkpoint_path, set_cpu=True)
 
     orig_image = Image.fromarray(input_image)
     image = resize_rmbg_image(orig_image)
@@ -213,7 +246,7 @@ def get_rmbg_matting(input_image: np.ndarray, checkpoint_path, ref_size=1024):
     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]
+    result = RMBG_SESS.run(None, {RMBG_SESS.get_inputs()[0].name: im_np})[0]
 
     # Post process
     result = np.squeeze(result)
@@ -271,3 +304,64 @@ def get_mnn_modnet_matting(input_image, checkpoint_path, ref_size=512):
     output_image = cv2.merge((b, g, r, mask))
 
     return output_image
+
+
+def get_birefnet_portrait_matting(input_image, checkpoint_path, ref_size=512):
+    global BIREFNET_V1_LITE_SESS
+
+    if not os.path.exists(checkpoint_path):
+        print(f"Checkpoint file not found: {checkpoint_path}")
+        return None
+
+    def transform_image(image):
+        image = image.resize((1024, 1024))  # Resize to 1024x1024
+        image = (
+            np.array(image, dtype=np.float32) / 255.0
+        )  # Convert to numpy array and normalize to [0, 1]
+        image = (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]  # Normalize
+        image = np.transpose(image, (2, 0, 1))  # Change from (H, W, C) to (C, H, W)
+        image = np.expand_dims(image, axis=0)  # Add batch dimension
+        return image.astype(np.float32)  # Ensure the output is float32
+
+    orig_image = Image.fromarray(input_image)
+    input_images = transform_image(
+        orig_image
+    )  # This will already have the correct shape
+
+    # 记录加载onnx模型的开始时间
+    load_start_time = time()
+
+    if BIREFNET_V1_LITE_SESS is None:
+        BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path, set_cpu=True)
+
+    # 记录加载onnx模型的结束时间
+    load_end_time = time()
+
+    # 打印加载onnx模型所花的时间
+    print(f"Loading ONNX model took {load_end_time - load_start_time:.4f} seconds")
+
+    input_name = BIREFNET_V1_LITE_SESS.get_inputs()[0].name
+    print(onnxruntime.get_device(), BIREFNET_V1_LITE_SESS.get_providers())
+
+    time_st = time()
+    pred_onnx = BIREFNET_V1_LITE_SESS.run(None, {input_name: input_images})[
+        -1
+    ]  # Use float32 input
+    pred_onnx = np.squeeze(pred_onnx)  # Use numpy to squeeze
+    result = 1 / (1 + np.exp(-pred_onnx))  # Sigmoid function using numpy
+    print(f"Inference time: {time() - time_st:.4f} seconds")
+
+    # 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)

hivision/creator/weights/birefnet-v1-lite.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5600024376f572a557870a5eb0afb1e5961636bef4e1e22132025467d0f03333
+size 224005088