update

demo/locales.py

@@ -620,4 +620,22 @@ LOCALES = {
             "label": "포화도 강도",
         },
     },
+    "face_alignment": {
+        "en": {
+            "label": "Face Alignment",
+            "value": "Open",
+        },
+        "zh": {
+            "label": "人脸旋转对齐",
+            "value": "开启",
+        },
+        "ja": {
+            "label": "顔の整列",
+            "value": "開始",
+        },
+        "ko": {
+            "label": "얼굴 정렬",
+            "value": "시작",
+        },
+    },
 }

demo/processor.py

@@ -53,6 +53,7 @@ class IDPhotoProcessor:
         contrast_strength=0,
         sharpen_strength=0,
         saturation_strength=0,
+        face_alignment_option=False,
     ):
         # 初始化参数
         top_distance_min = top_distance_max - 0.02
@@ -115,6 +116,7 @@ class IDPhotoProcessor:
                 contrast_strength,
                 sharpen_strength,
                 saturation_strength,
+                face_alignment_option,
             )
         except (FaceError, APIError):
             return self._handle_photo_generation_error(language)
@@ -209,6 +211,7 @@ class IDPhotoProcessor:
         contrast_strength,
         sharpen_strength,
         saturation_strength,
+        face_alignment_option,
     ):
         """生成证件照"""
         change_bg_only = (
@@ -225,6 +228,7 @@ class IDPhotoProcessor:
             contrast_strength=contrast_strength,
             sharpen_strength=sharpen_strength,
             saturation_strength=saturation_strength,
+            face_alignment=face_alignment_option,
         )
 
     def _handle_photo_generation_error(self, language):

demo/ui.py

@@ -23,7 +23,13 @@ def create_ui(
     face_detect_models: list,
     language: list,
 ):
-    DEFAULT_LANG = "en"
+
+    # 加载环境变量DEFAULT_LANG, 如果有且在language中，则将DEFAULT_LANG设置为环境变量
+    if "DEFAULT_LANG" in os.environ and os.environ["DEFAULT_LANG"] in language:
+        DEFAULT_LANG = os.environ["DEFAULT_LANG"]
+    else:
+        DEFAULT_LANG = language[0]
+
     DEFAULT_HUMAN_MATTING_MODEL = "modnet_photographic_portrait_matting"
     DEFAULT_FACE_DETECT_MODEL = "retinaface-resnet50"
 
@@ -67,12 +73,18 @@ def create_ui(
                 with gr.Tab(
                     LOCALES["key_param"][DEFAULT_LANG]["label"]
                 ) as key_parameter_tab:
-                    mode_options = gr.Radio(
-                        choices=LOCALES["size_mode"][DEFAULT_LANG]["choices"],
-                        label=LOCALES["size_mode"][DEFAULT_LANG]["label"],
-                        value=LOCALES["size_mode"][DEFAULT_LANG]["choices"][0],
-                    )
-
+                    with gr.Row():
+                        mode_options = gr.Radio(
+                            choices=LOCALES["size_mode"][DEFAULT_LANG]["choices"],
+                            label=LOCALES["size_mode"][DEFAULT_LANG]["label"],
+                            value=LOCALES["size_mode"][DEFAULT_LANG]["choices"][0],
+                            min_width=500,
+                        )
+                        face_alignment_options = gr.CheckboxGroup(
+                            label=LOCALES["face_alignment"][DEFAULT_LANG]["label"],
+                            choices=[LOCALES["face_alignment"][DEFAULT_LANG]["value"]],
+                            interactive=True,
+                        )
                     with gr.Row(visible=True) as size_list_row:
                         size_list_options = gr.Dropdown(
                             choices=LOCALES["size_list"][DEFAULT_LANG]["choices"],
@@ -80,7 +92,6 @@ def create_ui(
                             value=LOCALES["size_list"][DEFAULT_LANG]["choices"][0],
                             elem_id="size_list",
                         )
-
                     with gr.Row(visible=False) as custom_size:
                         custom_size_height = gr.Number(
                             value=413, label="height", interactive=True
@@ -483,6 +494,10 @@ def create_ui(
                     saturation_option: gr.update(
                         label=LOCALES["saturation_strength"][language]["label"]
                     ),
+                    face_alignment_options: gr.update(
+                        label=LOCALES["face_alignment"][language]["label"],
+                        value=LOCALES["face_alignment"][language]["value"],
+                    ),
                 }
 
             def change_visibility(option, lang, locales_key, custom_component):
@@ -497,20 +512,26 @@ def create_ui(
 
             def change_size_mode(size_option_item, lang):
                 choices = LOCALES["size_mode"][lang]["choices"]
+                # 如果选择自定义尺寸，则隐藏预设尺寸列表
                 if size_option_item == choices[2]:
                     return {
                         custom_size: gr.update(visible=True),
                         size_list_row: gr.update(visible=False),
+                        face_alignment_options: gr.update(visible=True),
                     }
+                # 如果选择只换底，则隐藏所有尺寸组件
                 elif size_option_item == choices[1]:
                     return {
                         custom_size: gr.update(visible=False),
                         size_list_row: gr.update(visible=False),
+                        face_alignment_options: gr.update(visible=False),
                     }
+                # 如果选择预设尺寸，则隐藏自定义尺寸组件
                 else:
                     return {
                         custom_size: gr.update(visible=False),
                         size_list_row: gr.update(visible=True),
+                        face_alignment_options: gr.update(visible=True),
                     }
 
             def change_image_kb(image_kb_option, lang):
@@ -566,6 +587,7 @@ def create_ui(
                     contrast_option,
                     sharpen_option,
                     saturation_option,
+                    face_alignment_options,
                 ],
             )
 
@@ -578,7 +600,7 @@ def create_ui(
             mode_options.input(
                 change_size_mode,
                 inputs=[mode_options, language_options],
-                outputs=[custom_size, size_list_row],
+                outputs=[custom_size, size_list_row, face_alignment_options],
             )
 
             image_kb_options.input(
@@ -627,6 +649,7 @@ def create_ui(
                     contrast_option,
                     sharpen_option,
                     saturation_option,
+                    face_alignment_options,
                 ],
                 outputs=[
                     img_output_standard,

hivision/creator/__init__.py

@@ -63,6 +63,7 @@ class IDCreator:
         contrast_strength: int = 0,
         sharpen_strength: int = 0,
         saturation_strength: int = 0,
+        face_alignment: bool = False,
     ) -> Result:
         """
         证件照处理函数
@@ -78,6 +79,8 @@ class IDCreator:
         :param brightness_strength: 亮度强度
         :param contrast_strength: 对比度强度
         :param sharpen_strength: 锐化强度
+        :param align_face: 是否需要人脸矫正
+
         :return: 返回处理后的证件照和一系列参数
         """
         # 0.初始化上下文
@@ -94,6 +97,7 @@ class IDCreator:
             contrast_strength=contrast_strength,
             sharpen_strength=sharpen_strength,
             saturation_strength=saturation_strength,
+            face_alignment=face_alignment,
         )
 
         self.ctx = Context(params)
@@ -105,7 +109,7 @@ class IDCreator:
         ctx.origin_image = ctx.processing_image.copy()
         self.before_all and self.before_all(ctx)
 
-        # 1. 人像抠图
+        # 1. ------------------人像抠图------------------
         if not ctx.params.crop_only:
             # 调用抠图工作流
             self.matting_handler(ctx)
@@ -113,7 +117,7 @@ class IDCreator:
         else:
             ctx.matting_image = ctx.processing_image
 
-        # 2. 美颜
+        # 2. ------------------美颜------------------
         self.beauty_handler(ctx)
 
         # 如果仅换底，则直接返回抠图结果
@@ -129,16 +133,33 @@ class IDCreator:
             self.after_all and self.after_all(ctx)
             return ctx.result
 
-        # 2. 人脸检测
+        # 3. ------------------人脸检测------------------
         self.detection_handler(ctx)
         self.after_detect and self.after_detect(ctx)
 
-        # 3. 图像调整
+        # 3.1 ------------------人脸对齐------------------
+        if ctx.params.face_alignment and abs(ctx.face["roll_angle"]) > 2:
+            from hivision.creator.rotation_adjust import rotate_bound_4channels
+
+            print("执行人脸对齐")
+            print("旋转角度：", ctx.face["roll_angle"])
+            # 根据角度旋转原图和抠图
+            ctx.origin_image, ctx.matting_image, _, _, _, _ = rotate_bound_4channels(
+                ctx.origin_image,
+                cv2.split(ctx.matting_image)[-1],
+                -1 * ctx.face["roll_angle"],
+            )
+
+            # 旋转后再执行一遍人脸检测
+            self.detection_handler(ctx)
+            self.after_detect and self.after_detect(ctx)
+
+        # 4. ------------------图像调整------------------
         result_image_hd, result_image_standard, clothing_params, typography_params = (
             adjust_photo(ctx)
         )
 
-        # 4. 返回结果
+        # 5. ------------------返回结果------------------
         ctx.result = Result(
             standard=result_image_standard,
             hd=result_image_hd,
@@ -148,4 +169,5 @@ class IDCreator:
             face=ctx.face,
         )
         self.after_all and self.after_all(ctx)
+
         return ctx.result

hivision/creator/context.py

@@ -26,6 +26,7 @@ class Params:
         contrast_strength: int = 0,
         sharpen_strength: int = 0,
         saturation_strength: int = 0,
+        face_alignment: bool = False,
     ):
         self.__size = size
         self.__change_bg_only = change_bg_only
@@ -39,6 +40,7 @@ class Params:
         self.__contrast_strength = contrast_strength
         self.__sharpen_strength = sharpen_strength
         self.__saturation_strength = saturation_strength
+        self.__face_alignment = face_alignment
 
     @property
     def size(self):
@@ -88,6 +90,10 @@ class Params:
     def saturation_strength(self):
         return self.__saturation_strength
 
+    @property
+    def face_alignment(self):
+        return self.__face_alignment
+
 
 class Result:
     def __init__(
@@ -143,15 +149,20 @@ class Context:
         """
         人像抠图结果
         """
-        self.face: Optional[Tuple[int, int, int, int, float]] = None
+        self.face: dict = dict(rectangle=None, roll_angle=None)
         """
         人脸检测结果，大于一个人脸时已在上层抛出异常
-        元组长度为5，包含 x1, y1, x2, y2, score 的坐标, (x1, y1)为左上角坐标，(x2, y2)为右下角坐标, score为置信度, 最大值为1
+        rectangle: 人脸矩形框，包含 x1, y1, width, height 的坐标, x1, y1 为左上角坐标, width, height 为矩形框的宽度和高度
+        roll_angle: 人脸偏转角度，以眼睛为标准，计算的人脸偏转角度，用于人脸矫正
         """
         self.result: Optional[Result] = None
         """
         证件照处理结果
         """
+        self.align_info: Optional[dict] = None
+        """
+        人脸矫正信息，仅当 align_face 为 True 时存在
+        """
 
 
 ContextHandler = Optional[Callable[[Context], None]]

hivision/creator/face_detector.py

@@ -20,6 +20,7 @@ from hivision.creator.retinaface import retinaface_detect_faces
 import requests
 import cv2
 import os
+import numpy as np
 
 
 mtcnn = None
@@ -42,11 +43,13 @@ def detect_face_mtcnn(ctx: Context, scale: int = 2):
         (ctx.origin_image.shape[1] // scale, ctx.origin_image.shape[0] // scale),
         interpolation=cv2.INTER_AREA,
     )
-    faces, _ = mtcnn.detect(image, thresholds=[0.8, 0.8, 0.8])
+    # landmarks 是 5 个关键点，分别是左眼、右眼、鼻子、左嘴角、右嘴角，
+    faces, landmarks = mtcnn.detect(image, thresholds=[0.8, 0.8, 0.8])
+
     # print(len(faces))
     if len(faces) != 1:
         # 保险措施，如果检测到多个人脸或者没有人脸，用原图再检测一次
-        faces, _ = mtcnn.detect(ctx.origin_image)
+        faces, landmarks = mtcnn.detect(ctx.origin_image)
     else:
         # 如果只有一个人脸，将人脸坐标放大
         for item, param in enumerate(faces[0]):
@@ -54,12 +57,23 @@ def detect_face_mtcnn(ctx: Context, scale: int = 2):
     if len(faces) != 1:
         raise FaceError("Expected 1 face, but got {}".format(len(faces)), len(faces))
 
+    # 计算人脸坐标
     left = faces[0][0]
     top = faces[0][1]
     width = faces[0][2] - left + 1
     height = faces[0][3] - top + 1
+    ctx.face["rectangle"] = (left, top, width, height)
+
+    # 根据landmarks计算人脸偏转角度，以眼睛为标准，计算的人脸偏转角度，用于人脸矫正
+    # 示例landmarks [106.37181  150.77415  127.21012  108.369156 144.61522  105.24723 107.45625  133.62355  151.24269  153.34407 ]
+    landmarks = landmarks[0]
+    left_eye = np.array([landmarks[0], landmarks[5]])
+    right_eye = np.array([landmarks[1], landmarks[6]])
+    dy = right_eye[1] - left_eye[1]
+    dx = right_eye[0] - left_eye[0]
+    roll_angle = np.degrees(np.arctan2(dy, dx))
 
-    ctx.face = (left, top, width, height)
+    ctx.face["roll_angle"] = roll_angle
 
 
 def detect_face_face_plusplus(ctx: Context):
@@ -83,6 +97,8 @@ def detect_face_face_plusplus(ctx: Context):
         "api_key": (None, api_key),
         "api_secret": (None, api_secret),
         "image_base64": (None, image_base64),
+        "return_landmark": (None, "1"),
+        "return_attributes": (None, "headpose"),
     }
 
     # 发送 POST 请求
@@ -96,12 +112,24 @@ def detect_face_face_plusplus(ctx: Context):
         face_num = response_json["face_num"]
         if face_num == 1:
             face_rectangle = response_json["faces"][0]["face_rectangle"]
-            ctx.face = (
+
+            # 获取人脸关键点
+            # landmarks = response_json["faces"][0]["landmark"]
+            # print("face++ landmarks", landmarks)
+
+            # headpose 是一个字典，包含俯仰角（pitch）、偏航角（yaw）和滚转角（roll）
+            # headpose示例 {'pitch_angle': 6.997899, 'roll_angle': 1.8011835, 'yaw_angle': 5.043002}
+            headpose = response_json["faces"][0]["attributes"]["headpose"]
+            # 以眼睛为标准，计算的人脸偏转角度，用于人脸矫正
+            roll_angle = headpose["roll_angle"] / 2
+
+            ctx.face["rectangle"] = (
                 face_rectangle["left"],
                 face_rectangle["top"],
                 face_rectangle["width"],
                 face_rectangle["height"],
             )
+            ctx.face["roll_angle"] = roll_angle
         else:
             raise FaceError(
                 "Expected 1 face, but got {}".format(face_num), len(face_num)
@@ -165,12 +193,26 @@ def detect_face_retinaface(ctx: Context):
         print("二次RetinaFace模型推理用时: {:.4f}s".format(time() - tic))
 
     faces_num = len(faces_dets)
+    faces_landmarks = []
+    for face_det in faces_dets:
+        faces_landmarks.append(face_det[5:])
+
     if faces_num != 1:
         raise FaceError("Expected 1 face, but got {}".format(faces_num), faces_num)
     face_det = faces_dets[0]
-    ctx.face = (
+    ctx.face["rectangle"] = (
         face_det[0],
         face_det[1],
         face_det[2] - face_det[0] + 1,
         face_det[3] - face_det[1] + 1,
     )
+
+    # 计算roll_angle
+    face_landmarks = faces_landmarks[0]
+    # print("face_landmarks", face_landmarks)
+    left_eye = np.array([face_landmarks[0], face_landmarks[1]])
+    right_eye = np.array([face_landmarks[2], face_landmarks[3]])
+    dy = right_eye[1] - left_eye[1]
+    dx = right_eye[0] - left_eye[0]
+    roll_angle = np.degrees(np.arctan2(dy, dx))
+    ctx.face["roll_angle"] = roll_angle

hivision/creator/photo_adjuster.py

@@ -17,7 +17,7 @@ import cv2
 
 def adjust_photo(ctx: Context):
     # Step1. 准备人脸参数
-    face_rect = ctx.face
+    face_rect = ctx.face["rectangle"]
     standard_size = ctx.params.size
     params = ctx.params
     x, y = face_rect[0], face_rect[1]

hivision/creator/rotation_adjust.py

@@ -0,0 +1,81 @@
+"""
+人脸旋转矫正模块
+
+本模块提供了用于旋转图像的函数，主要用于人脸旋转矫正。
+包含了处理3通道和4通道图像的旋转函数。
+"""
+
+import cv2
+import numpy as np
+
+
+def rotate_bound(image: np.ndarray, angle: float, center=None):
+    """
+    旋转图像而不损失信息的函数
+
+    Args:
+        image (np.ndarray): 输入图像，3通道numpy数组
+        angle (float): 旋转角度（度）
+        center (tuple, optional): 旋转中心坐标，默认为图像中心
+
+    Returns:
+        tuple: 包含以下元素的元组：
+            - rotated (np.ndarray): 旋转后的图像
+            - cos (float): 旋转角度的余弦值
+            - sin (float): 旋转角度的正弦值
+            - dW (int): 宽度变化量
+            - dH (int): 高度变化量
+    """
+    print("rotate_bound", image.shape)
+    (h, w) = image.shape[:2]
+    if center is None:
+        (cX, cY) = (w / 2, h / 2)
+    else:
+        (cX, cY) = center
+
+    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
+    cos = np.abs(M[0, 0])
+    sin = np.abs(M[0, 1])
+
+    nW = int((h * sin) + (w * cos))
+    nH = int((h * cos) + (w * sin))
+
+    M[0, 2] += (nW / 2) - cX
+    M[1, 2] += (nH / 2) - cY
+
+    rotated = cv2.warpAffine(image, M, (nW, nH))
+
+    # 计算偏移量
+    dW = nW - w
+    dH = nH - h
+
+    return rotated, cos, sin, dW, dH
+
+
+def rotate_bound_4channels(image: np.ndarray, a: np.ndarray, angle: float, center=None):
+    """
+    旋转4通道图像的函数
+
+    这是rotate_bound函数的4通道版本，可以同时处理RGB图像和其对应的alpha通道。
+
+    Args:
+        image (np.ndarray): 输入的3通道RGB图像
+        a (np.ndarray): 输入图像的alpha通道
+        angle (float): 旋转角度（度）
+        center (tuple, optional): 旋转中心坐标，默认为图像中心
+
+    Returns:
+        tuple: 包含以下元素的元组：
+            - input_image (np.ndarray): 旋转后的3通道RGB图像
+            - result_image (np.ndarray): 旋转后的4通道RGBA图像
+            - cos (float): 旋转角度的余弦值
+            - sin (float): 旋转角度的正弦值
+            - dW (int): 宽度变化量
+            - dH (int): 高度变化量
+    """
+    input_image, cos, sin, dW, dH = rotate_bound(image, angle, center)
+    new_a, _, _, _, _ = rotate_bound(a, angle, center)  # 对alpha通道进行旋转
+    b, g, r = cv2.split(input_image)
+    result_image = cv2.merge((b, g, r, new_a))  # 合并旋转后的RGB通道和alpha通道
+
+    return input_image, result_image, cos, sin, dW, dH

hivision/creator/utils.py

@@ -140,3 +140,67 @@ def detect_distance(value, crop_height, max=0.06, min=0.04):
         move_value = int(move_value * crop_height)
         # print("下移{}".format(move_value))
         return -1, move_value
+
+
+def cutting_rect_pan(
+    x1, y1, x2, y2, width, height, L1, L2, L3, clockwise, standard_size
+):
+    """
+    本函数的功能是对旋转矫正结果图的裁剪框进行修正 ———— 解决"旋转三角形"现象。
+    Args:
+        - x1: int, 裁剪框左上角的横坐标
+        - y1: int, 裁剪框左上角的纵坐标
+        - x2: int, 裁剪框右下角的横坐标
+        - y2: int, 裁剪框右下角的纵坐标
+        - width: int, 待裁剪图的宽度
+        - height:int, 待裁剪图的高度
+        - L1: CLassObject, 根据旋转点连线所构造函数
+        - L2: CLassObject, 根据旋转点连线所构造函数
+        - L3: ClassObject, 一个特殊裁切点的坐标
+        - clockwise: int, 旋转时针状态
+        - standard_size: tuple, 标准照的尺寸
+
+    Returns:
+        - x1: int, 新的裁剪框左上角的横坐标
+        - y1: int, 新的裁剪框左上角的纵坐标
+        - x2: int, 新的裁剪框右下角的横坐标
+        - y2: int, 新的裁剪框右下角的纵坐标
+        - x_bias: int, 裁剪框横坐标方向上的计算偏置量
+        - y_bias: int, 裁剪框纵坐标方向上的计算偏置量
+    """
+    # 用于计算的裁剪框坐标x1_cal,x2_cal,y1_cal,y2_cal(如果裁剪框超出了图像范围，则缩小直至在范围内)
+    x1_std = x1 if x1 > 0 else 0
+    x2_std = x2 if x2 < width else width
+    # y1_std = y1 if y1 > 0 else 0
+    y2_std = y2 if y2 < height else height
+
+    # 初始化x和y的计算偏置项x_bias和y_bias
+    x_bias = 0
+    y_bias = 0
+
+    # 如果顺时针偏转
+    if clockwise == 1:
+        if y2 > L1.forward_x(x1_std):
+            y_bias = int(-(y2_std - L1.forward_x(x1_std)))
+        if y2 > L2.forward_x(x2_std):
+            x_bias = int(-(x2_std - L2.forward_y(y2_std)))
+        x2 = x2_std + x_bias
+        if x1 < L3.x:
+            x1 = L3.x
+    # 如果逆时针偏转
+    else:
+        if y2 > L1.forward_x(x1_std):
+            x_bias = int(L1.forward_y(y2_std) - x1_std)
+        if y2 > L2.forward_x(x2_std):
+            y_bias = int(-(y2_std - L2.forward_x(x2_std)))
+        x1 = x1_std + x_bias
+        if x2 > L3.x:
+            x2 = L3.x
+
+    # 计算裁剪框的y的变化
+    y2 = int(y2_std + y_bias)
+    new_cut_width = x2 - x1
+    new_cut_height = int(new_cut_width / standard_size[1] * standard_size[0])
+    y1 = y2 - new_cut_height
+
+    return x1, y1, x2, y2, x_bias, y_bias