From 87d38480888e585f2bc0e75d37dc2752c10a7aa4 Mon Sep 17 00:00:00 2001
From: zhouchengrong <Zchengrong@yeah.net>
Date: Sun, 1 Dec 2024 19:49:27 +0800
Subject: [PATCH] =?UTF-8?q?feat=20=20flux=20=E5=8F=96=E6=B6=88=E6=B1=A1?=
 =?UTF-8?q?=E7=82=B9=E6=A3=80=E6=B5=8B=20=E5=A2=9E=E5=8A=A0=E7=B1=BB?=
 =?UTF-8?q?=E5=88=AB=E5=88=A4=E6=96=AD=20fix?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../generate_image/service_generate_image.py  | 41 +++++++++++--------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/app/service/generate_image/service_generate_image.py b/app/service/generate_image/service_generate_image.py
index daad5c6..2ba7862 100644
--- a/app/service/generate_image/service_generate_image.py
+++ b/app/service/generate_image/service_generate_image.py
@@ -92,23 +92,28 @@ class GenerateImage:
             image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
             is_smudge = True
             if self.category == "sketch" and self.version == "fast":
-                # 色阶调整
-                cutoff = 1
-                levels_img = autoLevels(image_result, cutoff)
-                # 亮度调整
-                luminance = luminance_adjust(0.3, levels_img)
-                # 去背景
-                remove_bg_image = remove_background(luminance)
-                # 人脸检测
-                # 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, 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)
-                image_result = not_smudge_image
+                if self.version == "fast":
+                    # 色阶调整
+                    cutoff = 1
+                    levels_img = autoLevels(image_result, cutoff)
+                    # 亮度调整
+                    luminance = luminance_adjust(0.3, levels_img)
+                    # 去背景
+                    remove_bg_image = remove_background(luminance)
+                    # 人脸检测
+                    # 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, 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)
+                    image_result = not_smudge_image
+                else:
+                    category, scores, not_smudge_image = generate_category_recognition(image=image_result, gender=self.gender)
+                    self.generate_data['category'] = str(category)
+                    image_result = not_smudge_image
             if is_smudge:  # 无污点
                 # image_result = adjust_contrast(image_result)
                 image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
@@ -190,7 +195,7 @@ def infer_cancel(tasks_id):
 if __name__ == '__main__':
     rd = GenerateImageModel(
         tasks_id="123-89",
-        prompt='a fabric print, flower, yellow, 4k, hud',
+        prompt='a single item of sketch of Wabi-sabi, skirt, tiered, 4k, white background',
         image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg",
         mode='txt2img',
         category="test",