feat 新增人脸识别 中心点检测

app/service/generate_image/service.py

@@ -23,7 +23,7 @@ from tritonclient.utils import np_to_triton_dtype
 from app.core.config import *
 from app.schemas.generate_image import GenerateImageModel
 from app.service.generate_image.utils.adjust_contrast import adjust_contrast
-from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust
+from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust, face_detect_pic
 from app.service.generate_image.utils.upload_sd_image import upload_png_sd
 
 logger = logging.getLogger()
@@ -85,19 +85,23 @@ class GenerateImage:
             image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
             is_smudge = True
             if self.category == "sketch":
-                # 色阶调整
+                # 人脸检测
-                cutoff = 1
+                if face_detect_pic(image_result) > 0:
-                levels_img = autoLevels(image_result, cutoff)
+                    is_smudge = False
-                # 亮度调整
+                else:
-                luminance = luminance_adjust(0.3, levels_img)
+                    # 色阶调整
-                # 去背景
+                    cutoff = 1
-                remove_bg_image = remove_background(luminance)
+                    levels_img = autoLevels(image_result, cutoff)
-                # 污点检测
+                    # 亮度调整
-                is_smudge, not_smudge_image = stain_detection(remove_bg_image)
+                    luminance = luminance_adjust(0.3, levels_img)
-                # 类型识别
+                    # 去背景
-                category, scores, not_smudge_image = generate_category_recognition(image=remove_bg_image, gender=self.gender)
+                    remove_bg_image = remove_background(luminance)
-                self.generate_data['category'] = str(category)
+                    # 污点/
-                image_result = not_smudge_image
+                    is_smudge, not_smudge_image = stain_detection(remove_bg_image)
+                    # 类型识别
+                    category, scores, not_smudge_image = generate_category_recognition(image=remove_bg_image, gender=self.gender)
+                    self.generate_data['category'] = str(category)
+                    image_result = not_smudge_image
             if is_smudge:  # 无污点
                 # image_result = adjust_contrast(image_result)
                 image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")

app/service/generate_image/utils/image_processing.py

@@ -154,11 +154,24 @@ def stain_detection(image, spot_size=100):
         if num_white_pixels != spot_size * spot_size:
             logger.info(f"第{index + 1}发现了污点")
             return False, None
-
+    # 中心区域检测
+    center_x, center_y = width // 2, height // 2
+    center_image = image[center_y - spot_size // 2:center_y + spot_size // 2, center_x - spot_size // 2:center_x + spot_size // 2]
+    gray_center_image = cv2.cvtColor(center_image, cv2.COLOR_BGR2GRAY)
+    white_threshold = 200  # 设定的接近最大值的阈值
+    is_pure_white = (gray_center_image >= white_threshold).all()
+    if is_pure_white:
+        return False, None
     if DEBUG:
-        for corner_coords in [(0, 0), (0, width - spot_size), (height - spot_size, 0), (height - spot_size, width - spot_size)]:
+        for corner_coords in [
+            (0, 0),
+            # (0, width - spot_size),
+            (height - spot_size, 0),
+            # (height - spot_size, width - spot_size)
+            # 中心点
+        ]:
             cv2.rectangle(image, corner_coords, (corner_coords[0] + spot_size, corner_coords[1] + spot_size), (0, 0, 255), 2)
-
+        cv2.rectangle(image, (center_x - spot_size // 2, center_y - spot_size // 2), (center_x + spot_size // 2, center_y + spot_size // 2), (0, 255, 0), 2)  # 在原始图像上绘制矩形框
     return True, image
 
 
@@ -241,6 +254,20 @@ def luminance_adjust(alpha, img):
     # 14.14 Photoshop 自动色阶调整算法
 
 
+def face_detect_pic(image):
+    # 1、转灰度图
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    # cv2.imshow("gray", gray)
+
+    # 2、训练一组人脸
+    face_detector = cv2.CascadeClassifier("service/generate_image/utils/haarcascade_frontalface_alt.xml")
+
+    # 3、检测人脸（用灰度图检测，返回人脸矩形坐标(4个角)）
+    faces_rect = face_detector.detectMultiScale(gray, 1.05, 3)
+
+    return len(faces_rect)
+
+
 if __name__ == '__main__':
     # Photoshop 自动色阶调整算法
     img = cv2.imread("2.png", flags=1)  # 读取彩色图像