feat generate 污点图保存

app/service/generate_image/service.py

@@ -24,7 +24,7 @@ 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, face_detect_pic
-from app.service.generate_image.utils.upload_sd_image import upload_png_sd
+from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_stain_png_sd
 
 logger = logging.getLogger()
 
@@ -93,11 +93,11 @@ class GenerateImage:
                 # 去背景
                 remove_bg_image = remove_background(luminance)
                 # 人脸检测
-                if face_detect_pic(remove_bg_image) > 0:
+                if face_detect_pic(remove_bg_image, self.user_id, self.category, self.tasks_id) > 0:
                     is_smudge = False
                 else:
                     # 污点/
-                    is_smudge, not_smudge_image = stain_detection(remove_bg_image)
+                    is_smudge, not_smudge_image = stain_detection(remove_bg_image, self.user_id, self.category, self.tasks_id)
                     # 类型识别
                     category, scores, not_smudge_image = generate_category_recognition(image=remove_bg_image, gender=self.gender)
                     self.generate_data['category'] = str(category)
@@ -111,7 +111,6 @@ class GenerateImage:
                 self.generate_data['image_url'] = str(image_url)
                 self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
             else:  # 有污点 保存图片到本地 测试用
-                cv2.imwrite(f"{self.tasks_id}.png", image_result)
                 self.generate_data['status'] = "SUCCESS"
                 self.generate_data['message'] = "success"
                 self.generate_data['image_url'] = str(GI_SYS_IMAGE_URL)

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.imshow("", dst)
-        cv2.waitKey(0)
+        # 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}")