Merge branch 'master' of https://github.com/Zeyi-Lin/HivisionIDPhotos

demo/processor.py

@@ -15,6 +15,10 @@
 
 
 class IDPhotoProcessor:
+    def __init__(self):
+        self.crop_only = False
+        self.matting_cache = None
+
     def process(
         self,
         input_image,
@@ -115,14 +119,46 @@ def process(
             idphoto_json["size_mode"] in LOCALES["size_mode"][language]["choices"][1]
         )
 
+        # 如果不是只裁切模式，则清空抠图缓存
+        if not self.crop_only:
+            # print("--清空缓存")
+            self.matting_cache = None
+            self.face_detect_cache = None
+
         try:
-            result = creator(
-                input_image,
-                change_bg_only=change_bg_only,
-                size=idphoto_json["size"],
-                head_measure_ratio=head_measure_ratio,
-                head_top_range=(top_distance_max, top_distance_min),
-            )
+            if self.matting_cache is None:
+                result = creator(
+                    input_image,
+                    change_bg_only=change_bg_only,
+                    size=idphoto_json["size"],
+                    head_measure_ratio=head_measure_ratio,
+                    head_top_range=(top_distance_max, top_distance_min),
+                )
+                self.matting_cache = result.matting
+                self.face_detect_cache = result.face
+            # 如果第一次点了只换底，第二次选了尺寸
+            elif self.matting_cache is not None and self.face_detect_cache is None:
+                result = creator(
+                    self.matting_cache,
+                    change_bg_only=change_bg_only,
+                    size=idphoto_json["size"],
+                    head_measure_ratio=head_measure_ratio,
+                    head_top_range=(top_distance_max, top_distance_min),
+                    crop_only=True,
+                )
+                self.face_detect_cache = result.face
+
+            else:
+                # print("使用缓存")
+                result = creator(
+                    self.matting_cache,
+                    change_bg_only=change_bg_only,
+                    size=idphoto_json["size"],
+                    head_measure_ratio=head_measure_ratio,
+                    head_top_range=(top_distance_max, top_distance_min),
+                    crop_only=True,
+                    face=self.face_detect_cache,
+                )
         except FaceError:
             return [
                 gr.update(value=None),  # img_output_standard
@@ -152,7 +188,7 @@ def process(
             ]
 
         else:
-            (result_image_standard, result_image_hd, _, _) = result
+            (result_image_standard, result_image_hd, _, _, _, _) = result
 
             result_image_standard_png = np.uint8(result_image_standard)
             result_image_hd_png = np.uint8(result_image_hd)
@@ -277,6 +313,9 @@ def process(
             else:
                 output_image_path = None
 
+            # 设置crop_only模式为True，会在Gradio的上传图像组件变化时重新设置为False
+            self.crop_only = True
+
             # 返回结果
             if output_image_path:
                 return [

demo/ui.py

@@ -2,7 +2,12 @@
 import os
 import pathlib
 from demo.locals import LOCALES
-from hivision.creator.choose_handler import FACE_DETECT_MODELS
+from demo.processor import IDPhotoProcessor
+
+"""
+只裁切模式:
+1. 如果重新上传了照片，然后点击按钮，第一次会调用不裁切的模式，第二次会调用裁切的模式
+"""
 
 
 def load_description(fp):
@@ -12,7 +17,10 @@ def load_description(fp):
 
 
 def create_ui(
-    processor, root_dir, human_matting_models: list, face_detect_models: list
+    processor: IDPhotoProcessor,
+    root_dir: str,
+    human_matting_models: list,
+    face_detect_models: list,
 ):
     DEFAULT_LANG = "zh"
     DEFAULT_HUMAN_MATTING_MODEL = "modnet_photographic_portrait_matting"
@@ -34,6 +42,10 @@ def create_ui(
             with gr.Column():
                 img_input = gr.Image(height=400)
 
+                def crop_only_false():
+                    # print("crop_only_false")
+                    processor.crop_only = False
+
                 with gr.Row():
                     # 语言选择器
                     language = ["zh", "en"]
@@ -55,6 +67,11 @@ def create_ui(
                         value=human_matting_models[0],
                     )
 
+                    # 重新上传照片、选择人脸检测模型、选择抠图模型时，重置裁切模式
+                    img_input.change(crop_only_false)
+                    face_detect_model_options.change(crop_only_false)
+                    matting_model_options.change(crop_only_false)
+
                 with gr.Tab(
                     LOCALES["key_param"][DEFAULT_LANG]["label"]
                 ) as key_parameter_tab:

hivision/creator/__init__.py

@@ -51,9 +51,11 @@ def __call__(
         image: np.ndarray,
         size: Tuple[int, int] = (413, 295),
         change_bg_only: bool = False,
+        crop_only: bool = False,
         head_measure_ratio: float = 0.2,
         head_height_ratio: float = 0.45,
         head_top_range: float = (0.12, 0.1),
+        face: Tuple[int, int, int, int] = None,
     ) -> Result:
         """
         证件照处理函数
@@ -73,6 +75,8 @@ def __call__(
             head_measure_ratio=head_measure_ratio,
             head_height_ratio=head_height_ratio,
             head_top_range=head_top_range,
+            crop_only=crop_only,
+            face=face,
         )
         self.ctx = Context(params)
         ctx = self.ctx
@@ -82,20 +86,32 @@ def __call__(
         )  # 将输入图片 resize 到最大边长为 2000
         ctx.origin_image = ctx.processing_image.copy()
         self.before_all and self.before_all(ctx)
+
         # 1. 人像抠图
-        self.matting_handler(ctx)
-        self.after_matting and self.after_matting(ctx)
+        if not ctx.params.crop_only:
+            self.matting_handler(ctx)
+            self.after_matting and self.after_matting(ctx)
+        else:
+            ctx.matting_image = ctx.processing_image
+
+        # 如果仅换底，则直接返回抠图结果
         if ctx.params.change_bg_only:
             ctx.result = Result(
                 standard=ctx.matting_image,
                 hd=ctx.matting_image,
+                matting=ctx.matting_image,
                 clothing_params=None,
                 typography_params=None,
+                face=None,
             )
             self.after_all and self.after_all(ctx)
             return ctx.result
+
         # 2. 人脸检测
-        self.detection_handler(ctx)
+        if not ctx.params.face:
+            self.detection_handler(ctx)
+        else:
+            ctx.face = ctx.params.face
         self.after_detect and self.after_detect(ctx)
         # 3. 图像调整
         result_image_hd, result_image_standard, clothing_params, typography_params = (
@@ -104,8 +120,10 @@ def __call__(
         ctx.result = Result(
             standard=result_image_standard,
             hd=result_image_hd,
+            matting=ctx.matting_image,
             clothing_params=clothing_params,
             typography_params=typography_params,
+            face=ctx.face,
         )
         self.after_all and self.after_all(ctx)
         return ctx.result

hivision/creator/context.py

@@ -16,15 +16,19 @@ def __init__(
         self,
         size: Tuple[int, int] = (413, 295),
         change_bg_only: bool = False,
+        crop_only: bool = False,
         head_measure_ratio: float = 0.2,
         head_height_ratio: float = 0.45,
         head_top_range: float = (0.12, 0.1),
+        face: Tuple[int, int, int, int] = None,
     ):
         self.__size = size
         self.__change_bg_only = change_bg_only
+        self.__crop_only = crop_only
         self.__head_measure_ratio = head_measure_ratio
         self.__head_height_ratio = head_height_ratio
         self.__head_top_range = head_top_range
+        self.__face = face
 
     @property
     def size(self):
@@ -46,17 +50,28 @@ def head_height_ratio(self):
     def head_top_range(self):
         return self.__head_top_range
 
+    @property
+    def crop_only(self):
+        return self.__crop_only
+
+    @property
+    def face(self):
+        return self.__face
+
 
 class Result:
     def __init__(
         self,
         standard: np.ndarray,
         hd: np.ndarray,
+        matting: np.ndarray,
         clothing_params: Optional[dict],
         typography_params: Optional[dict],
+        face: Optional[Tuple[int, int, int, int, float]],
     ):
         self.standard = standard
         self.hd = hd
+        self.matting = matting
         self.clothing_params = clothing_params
         """
         服装参数，仅换底时为 None
@@ -65,10 +80,18 @@ def __init__(
         """
         排版参数，仅换底时为 None
         """
+        self.face = face
 
     def __iter__(self):
         return iter(
-            [self.standard, self.hd, self.clothing_params, self.typography_params]
+            [
+                self.standard,
+                self.hd,
+                self.matting,
+                self.clothing_params,
+                self.typography_params,
+                self.face,
+            ]
         )