feat generate sketch 污点检测改为顶部两角

2024-04-25 15:54:17 +08:00
parent 78d0eb7373
commit a0c47c3367
1 changed files with 4 additions and 1 deletions
--- a/app/service/generate_image/service.py
+++ b/app/service/generate_image/service.py
@@ -78,7 +78,10 @@ class GenerateImage:
            self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
        else:
            image = result.as_numpy("generated_image")
-            image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_BGR2RGB)
+            if self.mode == "img2img":
+                image_result = np.squeeze(image.astype(np.uint8))
+            else:
+                image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_BGR2RGB)
            is_smudge = True
            if self.category == "sketch":
                # 去背景