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)