feat 新增 生成sketch时对图片清理背景，剔除带有污点的结果图

app/service/generate_image/utils/image_processing.py

@@ -1,11 +1,15 @@
+import logging
+
 import cv2
 import mmcv
 import numpy as np
 import torch
-from PIL import Image
 import tritonclient.http as httpclient
 import torch.nn.functional as F
 from app.core.config import *
+import cv2
+
+logger = logging.getLogger()
 
 
 def seg_preprocess(img_path):
@@ -107,11 +111,15 @@ def remove_background(image):
 
     result_mask = front_mask + back_mask
     white_background = np.ones_like(image_obj) * 255
-    result_image = np.where(result_mask[:, :, None].astype(bool), image_obj, white_background)
+    remove_bg_image = np.where(result_mask[:, :, None].astype(bool), image_obj, white_background)
+    # cv2.imwrite("source_image", image)
+    # cv2.imwrite("remove_bg_image", remove_bg_image)
 
-    import cv2
+    return remove_bg_image
 
-    edges = cv2.Canny(result_image, 50, 150)
+
+def bounding_box(image):
+    edges = cv2.Canny(image, 50, 150)
     # 查找轮廓
     contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     # 初始化包围所有外接矩形的大矩形的坐标
@@ -126,7 +134,29 @@ def remove_background(image):
 
     # 根据大矩形的坐标来裁剪原始图像
     result_image = image[y_min:y_max, x_min:x_max]
-    # cv2.imshow("", cropped_image)
+    # cv2.imshow("result_image", result_image)
     # cv2.waitKey(0)
-
     return result_image
+
+
+def stain_detection(image, spot_size=200):
+    height, width, _ = image.shape
+
+    corners = [
+        image[0:spot_size, 0:spot_size],  # top left
+        image[0:spot_size, width - spot_size:width],  # top right
+        image[height - spot_size:height, 0:spot_size],  # bottom left
+        image[height - spot_size:height, width - spot_size:width]  # bottom right
+    ]
+
+    for index, corner in enumerate(corners):
+        num_white_pixels = (corner == [255, 255, 255]).all(axis=2).sum()
+        if num_white_pixels != spot_size * spot_size:
+            logger.info(f"第{index + 1}发现了污点")
+            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)]:
+            cv2.rectangle(image, corner_coords, (corner_coords[0] + spot_size, corner_coords[1] + spot_size), (0, 0, 255), 2)
+
+    return True, image