update 20240924

demo/assets/title.md

@@ -1,7 +1,7 @@
 <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; text-align: center; font-size: 40px;">
   <div style="display: flex; align-items: center;">
     <img src="https://swanhub.co/git/repo/ZeYiLin%2FHivisionIDPhotos/file/preview?ref=master&path=assets/hivision_logo.png" alt="HivisionIDPhotos" style="width: 65px; height: 65px; margin-right: 10px;" onerror="this.style.display='none';" loading="lazy">
-    <b style="color: #6e9abb;">HivisionIDPhotos</b>
+    <b style="color: #6e9abb;">HivisionIDPhotos</b><span style="font-size: 18px; color: #638fb3; margin-left: 10px;"> v1.2.8</span>
   </div>
   <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
       <a href="https://github.com/xiaolin199912/HivisionIDPhotos"><img alt="Github" src="https://img.shields.io/static/v1?label=GitHub&message=GitHub&color=black"></a> &ensp;

demo/locales.py

@@ -667,22 +667,22 @@ LOCALES = {
             "label": "포화도 강도",
         },
     },
-    "face_alignment": {
+    "plugin": {
         "en": {
-            "label": "Face Alignment",
-            "choices": ["Open"],
+            "label": "🤖Plugin",
+            "choices": ["Face Alignment", "Layout Photo Crop Line"],
         },
         "zh": {
-            "label": "人脸旋转对齐",
-            "choices": ["开启"],
+            "label": "🤖插件",
+            "choices": ["人脸旋转对齐", "排版照裁剪线"],
         },
         "ja": {
-            "label": "顔の整列",
-            "choices": ["開始"],
+            "label": "🤖プラグイン",
+            "choices": ["顔の整列", "レイアウト写真の切り取り線"],
         },
         "ko": {
-            "label": "얼굴 정렬",
-            "choices": ["시작"],
+            "label": "🤖플러그인",
+            "choices": ["얼굴 정렬", "레이아웃 사진 자르기 선"],
         },
     },
     "template_photo": {

demo/processor.py

@@ -9,7 +9,7 @@ from hivision.utils import (
     save_image_dpi_to_bytes,
 )
 from hivision.creator.layout_calculator import (
-    generate_layout_photo,
+    generate_layout_array,
     generate_layout_image,
 )
 from hivision.creator.choose_handler import choose_handler
@@ -61,16 +61,26 @@ class IDPhotoProcessor:
         contrast_strength=0,
         sharpen_strength=0,
         saturation_strength=0,
-        face_alignment_option=False,
-    ):
+        plugin_option=[],
+    ):        
         # 初始化参数
         top_distance_min = top_distance_max - 0.02
         # 得到render_option在LOCALES["render_mode"][language]["choices"]中的索引
         render_option_index = LOCALES["render_mode"][language]["choices"].index(
             render_option
         )
+        # 读取插件选项
+        if LOCALES["plugin"][language]["choices"][0] in plugin_option:
+            face_alignment_option = True
+        else:
+            face_alignment_option = False
+        if LOCALES["plugin"][language]["choices"][1] in plugin_option:
+            layout_photo_crop_line_option = True
+        else:
+            layout_photo_crop_line_option = False
+        
         idphoto_json = self._initialize_idphoto_json(
-            mode_option, color_option, render_option_index, image_kb_options
+            mode_option, color_option, render_option_index, image_kb_options, layout_photo_crop_line_option
         )
 
         # 处理尺寸模式
@@ -153,6 +163,7 @@ class IDPhotoProcessor:
         color_option,
         render_option,
         image_kb_options,
+        layout_photo_crop_line_option,
     ):
         """初始化idphoto_json字典"""
         return {
@@ -162,6 +173,7 @@ class IDPhotoProcessor:
             "image_kb_mode": image_kb_options,
             "custom_image_kb": None,
             "custom_image_dpi": None,
+            "layout_photo_crop_line_option": layout_photo_crop_line_option,
         }
 
     # 处理尺寸模式
@@ -435,7 +447,7 @@ class IDPhotoProcessor:
         if idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]:
             return None, False
 
-        typography_arr, typography_rotate = generate_layout_photo(
+        typography_arr, typography_rotate = generate_layout_array(
             input_height=idphoto_json["size"][0],
             input_width=idphoto_json["size"][1],
         )
@@ -446,6 +458,7 @@ class IDPhotoProcessor:
             typography_rotate,
             height=idphoto_json["size"][0],
             width=idphoto_json["size"][1],
+            crop_line=idphoto_json["layout_photo_crop_line_option"],
         )
 
         return result_image_layout, True

demo/ui.py

@@ -81,12 +81,7 @@ def create_ui(
                             value=LOCALES["size_mode"][DEFAULT_LANG]["choices"][0],
                             min_width=520,
                         )
-                        # 人脸对齐
-                        face_alignment_options = gr.CheckboxGroup(
-                            label=LOCALES["face_alignment"][DEFAULT_LANG]["label"],
-                            choices=LOCALES["face_alignment"][DEFAULT_LANG]["choices"],
-                            interactive=True,
-                        )
+                        
                     # 尺寸列表
                     with gr.Row(visible=True) as size_list_row:
                         size_list_options = gr.Dropdown(
@@ -143,6 +138,14 @@ def create_ui(
                         label=LOCALES["render_mode"][DEFAULT_LANG]["label"],
                         value=LOCALES["render_mode"][DEFAULT_LANG]["choices"][0],
                     )
+                    
+                    with gr.Row():
+                        # 插件模式
+                        plugin_options = gr.CheckboxGroup(
+                            label=LOCALES["plugin"][DEFAULT_LANG]["label"],
+                            choices=LOCALES["plugin"][DEFAULT_LANG]["choices"],
+                            interactive=True,
+                        )
 
                 # TAB2 - 高级参数 ------------------------------------------------
                 with gr.Tab(
@@ -525,10 +528,6 @@ def create_ui(
                     saturation_option: gr.update(
                         label=LOCALES["saturation_strength"][language]["label"]
                     ),
-                    face_alignment_options: gr.update(
-                        label=LOCALES["face_alignment"][language]["label"],
-                        choices=LOCALES["face_alignment"][language]["choices"],
-                    ),
                     custom_size_width_px: gr.update(
                         label=LOCALES["custom_size_px"][language]["width"]
                     ),
@@ -547,6 +546,11 @@ def create_ui(
                     template_image_accordion: gr.update(
                         label=LOCALES["template_photo"][language]["label"]
                     ),
+                    plugin_options: gr.update(
+                        label=LOCALES["plugin"][language]["label"],
+                        choices=LOCALES["plugin"][language]["choices"],
+                        value=LOCALES["plugin"][language]["choices"][0],
+                    ),
                 }
 
             def change_visibility(option, lang, locales_key, custom_component):
@@ -571,7 +575,7 @@ def create_ui(
                         custom_size_px: gr.update(visible=False),
                         custom_size_mm: gr.update(visible=True),
                         size_list_row: gr.update(visible=False),
-                        face_alignment_options: gr.update(visible=True),
+                        plugin_options: gr.update(interactive=True),
                     }
                 # 如果选择自定义尺寸px
                 elif size_option_item == choices[2]:
@@ -579,7 +583,7 @@ def create_ui(
                         custom_size_px: gr.update(visible=True),
                         custom_size_mm: gr.update(visible=False),
                         size_list_row: gr.update(visible=False),
-                        face_alignment_options: gr.update(visible=True),
+                        plugin_options: gr.update(interactive=True),
                     }
                 # 如果选择只换底，则隐藏所有尺寸组件
                 elif size_option_item == choices[1]:
@@ -587,7 +591,7 @@ def create_ui(
                         custom_size_px: gr.update(visible=False),
                         custom_size_mm: gr.update(visible=False),
                         size_list_row: gr.update(visible=False),
-                        face_alignment_options: gr.update(visible=False),
+                        plugin_options: gr.update(interactive=False),
                     }
                 # 如果选择预设尺寸，则隐藏自定义尺寸组件
                 else:
@@ -595,7 +599,7 @@ def create_ui(
                         custom_size_px: gr.update(visible=False),
                         custom_size_mm: gr.update(visible=False),
                         size_list_row: gr.update(visible=True),
-                        face_alignment_options: gr.update(visible=True),
+                        plugin_options: gr.update(interactive=True),
                     }
 
             def change_image_kb(image_kb_option, lang):
@@ -650,7 +654,7 @@ def create_ui(
                     contrast_option,
                     sharpen_option,
                     saturation_option,
-                    face_alignment_options,
+                    plugin_options,
                     custom_size_width_px,
                     custom_size_height_px,
                     custom_size_width_mm,
@@ -669,7 +673,7 @@ def create_ui(
                     custom_size_px,
                     custom_size_mm,
                     size_list_row,
-                    face_alignment_options,
+                    plugin_options,
                 ],
             )
 
@@ -731,7 +735,7 @@ def create_ui(
                     contrast_option,
                     sharpen_option,
                     saturation_option,
-                    face_alignment_options,
+                    plugin_options,
                 ],
                 outputs=[
                     img_output_standard,

hivision/creator/__init__.py

@@ -16,6 +16,7 @@ from .face_detector import detect_face_mtcnn
 from hivision.plugin.beauty.handler import beauty_face
 from .photo_adjuster import adjust_photo
 import cv2
+import time
 
 
 class IDCreator:
@@ -100,6 +101,10 @@ class IDCreator:
             face_alignment=face_alignment,
         )
 
+
+        # 总的开始时间
+        total_start_time = time.time()
+        
         self.ctx = Context(params)
         ctx = self.ctx
         ctx.processing_image = image
@@ -110,15 +115,26 @@ class IDCreator:
         self.before_all and self.before_all(ctx)
 
         # 1. ------------------人像抠图------------------
+        # 如果仅裁剪，则不进行抠图
         if not ctx.params.crop_only:
             # 调用抠图工作流
+            print("[1]  Start Human Matting...")
+            start_matting_time = time.time()
             self.matting_handler(ctx)
+            end_matting_time = time.time()
+            print(f"[1]  Human Matting Time: {end_matting_time - start_matting_time:.3f}s")
             self.after_matting and self.after_matting(ctx)
+        # 如果进行抠图
         else:
             ctx.matting_image = ctx.processing_image
 
+
         # 2. ------------------美颜------------------
+        print("[2]  Start Beauty...")
+        start_beauty_time = time.time()
         self.beauty_handler(ctx)
+        end_beauty_time = time.time()
+        print(f"[2]  Beauty Time: {end_beauty_time - start_beauty_time:.3f}s")
 
         # 如果仅换底，则直接返回抠图结果
         if ctx.params.change_bg_only:
@@ -134,15 +150,19 @@ class IDCreator:
             return ctx.result
 
         # 3. ------------------人脸检测------------------
+        print("[3]  Start Face Detection...")
+        start_detection_time = time.time()
         self.detection_handler(ctx)
+        end_detection_time = time.time()
+        print(f"[3]  Face Detection Time: {end_detection_time - start_detection_time:.3f}s")
         self.after_detect and self.after_detect(ctx)
 
         # 3.1 ------------------人脸对齐------------------
         if ctx.params.face_alignment and abs(ctx.face["roll_angle"]) > 2:
+            print("[3.1]  Start Face Alignment...")
+            start_alignment_time = time.time()
             from hivision.creator.rotation_adjust import rotate_bound_4channels
 
-            print("执行人脸对齐")
-            print("旋转角度：", ctx.face["roll_angle"])
             # 根据角度旋转原图和抠图
             b, g, r, a = cv2.split(ctx.matting_image)
             ctx.origin_image, ctx.matting_image, _, _, _, _ = rotate_bound_4channels(
@@ -154,11 +174,17 @@ class IDCreator:
             # 旋转后再执行一遍人脸检测
             self.detection_handler(ctx)
             self.after_detect and self.after_detect(ctx)
+            end_alignment_time = time.time()
+            print(f"[3.1]  Face Alignment Time: {end_alignment_time - start_alignment_time:.3f}s")
 
         # 4. ------------------图像调整------------------
+        print("[4]  Start Image Post-Adjustment...")
+        start_adjust_time = time.time()
         result_image_hd, result_image_standard, clothing_params, typography_params = (
             adjust_photo(ctx)
         )
+        end_adjust_time = time.time()
+        print(f"[4]  Image Post-Adjustment Time: {end_adjust_time - start_adjust_time:.3f}s")
 
         # 5. ------------------返回结果------------------
         ctx.result = Result(
@@ -171,4 +197,8 @@ class IDCreator:
         )
         self.after_all and self.after_all(ctx)
 
+        # 总的结束时间
+        total_end_time = time.time()
+        print(f"[Total]  Total Time: {total_end_time - total_start_time:.3f}s")
+
         return ctx.result

hivision/creator/face_detector.py

@@ -173,7 +173,6 @@ def detect_face_retinaface(ctx: Context):
     global RETINAFCE_SESS
 
     if RETINAFCE_SESS is None:
-        print("首次加载RetinaFace模型...")
         # 计算用时
         tic = time()
         faces_dets, sess = retinaface_detect_faces(
@@ -182,7 +181,6 @@ def detect_face_retinaface(ctx: Context):
             sess=None,
         )
         RETINAFCE_SESS = sess
-        print("首次RetinaFace模型推理用时: {:.4f}s".format(time() - tic))
     else:
         tic = time()
         faces_dets, _ = retinaface_detect_faces(
@@ -190,7 +188,6 @@ def detect_face_retinaface(ctx: Context):
             os.path.join(base_dir, "retinaface/weights/retinaface-resnet50.onnx"),
             sess=RETINAFCE_SESS,
         )
-        print("二次RetinaFace模型推理用时: {:.4f}s".format(time() - tic))
 
     faces_num = len(faces_dets)
     faces_landmarks = []
@@ -216,3 +213,7 @@ def detect_face_retinaface(ctx: Context):
     dx = right_eye[0] - left_eye[0]
     roll_angle = np.degrees(np.arctan2(dy, dx))
     ctx.face["roll_angle"] = roll_angle
+
+    # 如果RUN_MODE不是野兽模式，则释放模型
+    if os.getenv("RUN_MODE") == "beast":
+        RETINAFCE_SESS = None

hivision/creator/human_matting.py

@@ -201,6 +201,7 @@ def get_modnet_matting(input_image, checkpoint_path, ref_size=512):
         print(f"Checkpoint file not found: {checkpoint_path}")
         return None
 
+    # 如果RUN_MODE不是野兽模式，则不加载模型
     if HIVISION_MODNET_SESS is None:
         HIVISION_MODNET_SESS = load_onnx_model(checkpoint_path, set_cpu=True)
 
@@ -216,6 +217,10 @@ def get_modnet_matting(input_image, checkpoint_path, ref_size=512):
     b, g, r = cv2.split(np.uint8(input_image))
 
     output_image = cv2.merge((b, g, r, mask))
+    
+    # 如果RUN_MODE不是野兽模式，则释放模型
+    if os.getenv("RUN_MODE") != "beast":
+        HIVISION_MODNET_SESS = None
 
     return output_image
 
@@ -229,6 +234,7 @@ def get_modnet_matting_photographic_portrait_matting(
         print(f"Checkpoint file not found: {checkpoint_path}")
         return None
 
+    # 如果RUN_MODE不是野兽模式，则不加载模型
     if MODNET_PHOTOGRAPHIC_PORTRAIT_MATTING_SESS is None:
         MODNET_PHOTOGRAPHIC_PORTRAIT_MATTING_SESS = load_onnx_model(
             checkpoint_path, set_cpu=True
@@ -248,6 +254,10 @@ def get_modnet_matting_photographic_portrait_matting(
     b, g, r = cv2.split(np.uint8(input_image))
 
     output_image = cv2.merge((b, g, r, mask))
+    
+    # 如果RUN_MODE不是野兽模式，则释放模型
+    if os.getenv("RUN_MODE") != "beast":
+        MODNET_PHOTOGRAPHIC_PORTRAIT_MATTING_SESS = None
 
     return output_image
 
@@ -297,6 +307,10 @@ def get_rmbg_matting(input_image: np.ndarray, checkpoint_path, ref_size=1024):
     # 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)
+    
+    # 如果RUN_MODE不是野兽模式，则释放模型
+    if os.getenv("RUN_MODE") != "beast":
+        RMBG_SESS = None
 
     return np.array(new_im)
 
@@ -362,8 +376,9 @@ def get_birefnet_portrait_matting(input_image, checkpoint_path, ref_size=512):
     # 记录加载onnx模型的开始时间
     load_start_time = time()
 
+    # 如果RUN_MODE不是野兽模式，则不加载模型
     if BIREFNET_V1_LITE_SESS is None:
-        print("首次加载birefnet-v1-lite模型...")
+        # print("首次加载birefnet-v1-lite模型...")
         if ONNX_DEVICE == "GPU":
             print("onnxruntime-gpu已安装，尝试使用CUDA加载模型")
             try:
@@ -405,5 +420,9 @@ def get_birefnet_portrait_matting(input_image, checkpoint_path, ref_size=512):
     # 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)
+    
+    # 如果RUN_MODE不是野兽模式，则释放模型
+    if os.getenv("RUN_MODE") != "beast":
+        BIREFNET_V1_LITE_SESS = None
 
     return np.array(new_im)

hivision/creator/layout_calculator.py

@@ -75,7 +75,7 @@ def judge_layout(
         return layout_mode, centerBlockWidth_1, centerBlockHeight_1
 
 
-def generate_layout_photo(input_height, input_width):
+def generate_layout_array(input_height, input_width):
     # 1.基础参数表
     LAYOUT_WIDTH = 1746
     LAYOUT_HEIGHT = 1180
@@ -118,23 +118,63 @@ def generate_layout_photo(input_height, input_width):
 
 
 def generate_layout_image(
-    input_image, typography_arr, typography_rotate, width=295, height=413
+    input_image, typography_arr, typography_rotate, width=295, height=413, crop_line:bool=False,
 ):
+    # 定义画布的宽度和高度
     LAYOUT_WIDTH = 1746
     LAYOUT_HEIGHT = 1180
+    
+    # 创建一个白色背景的空白画布
     white_background = np.zeros([LAYOUT_HEIGHT, LAYOUT_WIDTH, 3], np.uint8)
     white_background.fill(255)
+    
+    # 如果输入图像的高度不等于指定高度，则调整图像大小
     if input_image.shape[0] != height:
         input_image = cv2.resize(input_image, (width, height))
+    
+    # 如果需要旋转排版，则对图像进行转置和垂直镜像
     if typography_rotate:
         input_image = cv2.transpose(input_image)
         input_image = cv2.flip(input_image, 0)  # 0 表示垂直镜像
 
+        # 交换高度和宽度
         height, width = width, height
+    
+    # 将图像按照排版数组中的位置放置到白色背景上
     for arr in typography_arr:
         locate_x, locate_y = arr[0], arr[1]
         white_background[locate_y : locate_y + height, locate_x : locate_x + width] = (
             input_image
         )
 
+    if crop_line:
+        # 添加裁剪线
+        line_color = (200, 200, 200)  # 浅灰色
+        line_thickness = 1
+
+        # 初始化裁剪线位置列表
+        vertical_lines = []
+        horizontal_lines = []
+
+        # 根据排版数组添加裁剪线
+        for arr in typography_arr:
+            x, y = arr[0], arr[1]
+            if x not in vertical_lines:
+                vertical_lines.append(x)
+            if x + width not in vertical_lines:
+                vertical_lines.append(x + width)
+            if y not in horizontal_lines:
+                horizontal_lines.append(y)
+            if y + height not in horizontal_lines:
+                horizontal_lines.append(y + height)
+
+        # 绘制垂直裁剪线
+        for x in vertical_lines:
+            cv2.line(white_background, (x, 0), (x, LAYOUT_HEIGHT), line_color, line_thickness)
+
+        # 绘制水平裁剪线
+        for y in horizontal_lines:
+            cv2.line(white_background, (0, y), (LAYOUT_WIDTH, y), line_color, line_thickness)
+
+    # 返回排版后的图像
     return white_background

hivision/creator/photo_adjuster.py

@@ -8,7 +8,7 @@ r"""
     证件照调整
 """
 from .context import Context
-from .layout_calculator import generate_layout_photo
+from .layout_calculator import generate_layout_array
 import hivision.creator.utils as U
 import numpy as np
 import math
@@ -111,7 +111,7 @@ def adjust_photo(ctx: Context):
     }
 
     # Step7. 排版照参数获取
-    typography_arr, typography_rotate = generate_layout_photo(
+    typography_arr, typography_rotate = generate_layout_array(
         input_height=standard_size[0], input_width=standard_size[1]
     )
 
@@ -173,7 +173,6 @@ def IDphotos_cut(x1, y1, x2, y2, img):
         temp_x_2 = temp_x_2 - x2
 
     # 生成一张全透明背景
-    print("crop_size:", crop_size)
     background_bgr = np.full((crop_size[0], crop_size[1]), 255, dtype=np.uint8)
     background_a = np.full((crop_size[0], crop_size[1]), 0, dtype=np.uint8)
     background = cv2.merge(

hivision/creator/rotation_adjust.py

@@ -26,7 +26,6 @@ def rotate_bound(image: np.ndarray, angle: float, center=None):
             - dW (int): 宽度变化量
             - dH (int): 高度变化量
     """
-    print("rotate_bound", image.shape)
     (h, w) = image.shape[:2]
     if center is None:
         (cX, cY) = (w / 2, h / 2)

hivision/utils.py

@@ -8,7 +8,7 @@ import base64
 from hivision.plugin.watermark import Watermarker, WatermarkerStyles
 
 
-def save_image_dpi_to_bytes(image, output_image_path, dpi=300):
+def save_image_dpi_to_bytes(image: np.ndarray, output_image_path: str = None, dpi: int = 300):
     """
     设置图像的DPI（每英寸点数）并返回字节流
 
@@ -25,13 +25,14 @@ def save_image_dpi_to_bytes(image, output_image_path, dpi=300):
     image_bytes = byte_stream.getvalue()
 
     # Save the image to the output path
-    with open(output_image_path, "wb") as f:
-        f.write(image_bytes)
+    if output_image_path:
+        with open(output_image_path, "wb") as f:
+            f.write(image_bytes)
 
     return image_bytes
 
 
-def resize_image_to_kb(input_image, output_image_path, target_size_kb, dpi=300):
+def resize_image_to_kb(input_image: np.ndarray, output_image_path: str = None, target_size_kb: int = 100, dpi: int = 300):
     """
     Resize an image to a target size in KB.
     将图像调整大小至目标文件大小（KB）。
@@ -79,9 +80,11 @@ def resize_image_to_kb(input_image, output_image_path, target_size_kb, dpi=300):
                 img_byte_arr.write(padding)
 
             # Save the image to the output path
-            with open(output_image_path, "wb") as f:
-                f.write(img_byte_arr.getvalue())
-            break
+            if output_image_path:
+                with open(output_image_path, "wb") as f:
+                    f.write(img_byte_arr.getvalue())
+            
+            return img_byte_arr.getvalue()
 
         # Reduce the quality if the image is still too large
         quality -= 5
@@ -172,11 +175,26 @@ def numpy_2_base64(img: np.ndarray) -> str:
 
 
 def base64_2_numpy(base64_image: str) -> np.ndarray:
-    img = base64.b64decode(base64_image)
-    img = np.frombuffer(img, np.uint8)
-
+    # Remove the data URL prefix if present
+    if base64_image.startswith('data:image'):
+        base64_image = base64_image.split(',')[1]
+    
+    # Decode base64 string to bytes
+    img_bytes = base64.b64decode(base64_image)
+    
+    # Convert bytes to numpy array
+    img_array = np.frombuffer(img_bytes, dtype=np.uint8)
+    
+    # Decode the image array
+    img = cv2.imdecode(img_array, cv2.IMREAD_UNCHANGED)
+    
     return img
 
+# 字节流转base64
+def bytes_2_base64(img_byte_arr: bytes) -> str:
+    base64_image = base64.b64encode(img_byte_arr).decode("utf-8")
+    return "data:image/png;base64," + base64_image
+
 
 def save_numpy_image(numpy_img, file_path):
     # 检查数组的形状