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)
 
 

app/service/generate_image/utils/upload_sd_image.py

@@ -33,3 +33,35 @@ def upload_png_sd(image, user_id, category, object_name):
         return image_url
     except Exception as e:
         logging.warning(f"upload_png_mask runtime exception : {e}")
+
+
+def upload_stain_png_sd(image, user_id, category, object_name):
+    try:
+        _, img_byte_array = cv2.imencode('.jpg', image)
+        minio_req = minio_client.put_object(
+            "test",
+            f'generate_result/stain/{user_id}_{category}_{object_name}',
+            io.BytesIO(img_byte_array),
+            len(img_byte_array),
+            content_type='image/jpeg'
+        )
+        image_url = f"test/{minio_req.object_name}"
+        return image_url
+    except Exception as e:
+        logging.warning(f"upload_png_mask runtime exception : {e}")
+
+
+def upload_face_png_sd(image, user_id, category, object_name):
+    try:
+        _, img_byte_array = cv2.imencode('.jpg', image)
+        minio_req = minio_client.put_object(
+            "test",
+            f'generate_result/face/{user_id}_{category}_{object_name}',
+            io.BytesIO(img_byte_array),
+            len(img_byte_array),
+            content_type='image/jpeg'
+        )
+        image_url = f"test/{minio_req.object_name}"
+        return image_url
+    except Exception as e:
+        logging.warning(f"upload_png_mask runtime exception : {e}")