feat generate 污点图保存

app/service/generate_image/utils/image_processing.py

@@ -7,6 +7,8 @@ import torch.nn.functional as F
 from app.core.config import *
 import cv2
 
+from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_stain_png_sd, upload_face_png_sd
+
 logger = logging.getLogger()
 
 
@@ -137,7 +139,7 @@ def bounding_box(image):
     return result_image
 
 
-def stain_detection(image, spot_size=100):
+def stain_detection(image, user_id, category, tasks_id, spot_size=100):
     height, width, _ = image.shape
 
     corners = [
@@ -197,6 +199,17 @@ def stain_detection(image, spot_size=100):
         ]:
             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)  # 在原始图像上绘制矩形框
+    dst = image.copy()
+    for corner_coords in [
+        (0, 0),
+        # (0, width - spot_size),
+        (height - spot_size, 0),
+        # (height - spot_size, width - spot_size)
+        # 中心点
+    ]:
+        cv2.rectangle(dst, corner_coords, (corner_coords[0] + spot_size, corner_coords[1] + spot_size), (0, 0, 255), 2)
+    cv2.rectangle(dst, (center_x - spot_size // 2, center_y - spot_size // 2), (center_x + spot_size // 2, center_y + spot_size // 2), (0, 255, 0), 2)  # 在原始图像上绘制矩形框
+    image_url = upload_stain_png_sd(dst, user_id=user_id, category=f"{category}", object_name=f"{tasks_id}.png")
     return True, image
 
 
@@ -279,7 +292,7 @@ def luminance_adjust(alpha, img):
     # 14.14 Photoshop 自动色阶调整算法
 
 
-def face_detect_pic(image):
+def face_detect_pic(image, user_id, category, tasks_id):
     # 1、转灰度图
     gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
     # cv2.imshow("gray", gray)
@@ -294,8 +307,13 @@ def face_detect_pic(image):
         dst = image.copy()
         for x, y, w, h in faces_rect:
             cv2.rectangle(dst, (x, y), (x + w, y + h), (0, 0, 255), 3)  # 画出矩形框
-        cv2.imshow("", dst)
-        cv2.waitKey(0)
+        # cv2.imshow("", dst)
+        # cv2.waitKey(0)
+    # TODO 暂时保留
+    dst = image.copy()
+    for x, y, w, h in faces_rect:
+        cv2.rectangle(dst, (x, y), (x + w, y + h), (0, 0, 255), 3)  # 画出矩形框
+    image_url = upload_face_png_sd(dst, user_id=user_id, category=f"{category}", object_name=f"{tasks_id}.png")
     return len(faces_rect)