feat generate 新增背景去除 bounding box

app/service/generate_image/service.py

@@ -22,6 +22,7 @@ from tritonclient.utils import np_to_triton_dtype
 
 from app.core.config import *
 from app.schemas.generate_image import GenerateImageModel
+from app.service.generate_image.utils.remove_background import remove_background
 from app.service.generate_image.utils.upload_sd_image import upload_png_sd
 
 logger = logging.getLogger()
@@ -71,6 +72,8 @@ class GenerateImage:
             self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
         else:
             image_result = result.as_numpy("generated_image")[0]
+            if self.category == "sketch":
+                image_result = remove_background(np.asarray(image_result))
             image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")
             self.generate_data['status'] = "SUCCESS"
             self.generate_data['message'] = "success"

app/service/generate_image/utils/remove_background.py

@@ -109,4 +109,24 @@ def remove_background(image):
     white_background = np.ones_like(image_obj) * 255
     result_image = np.where(result_mask[:, :, None].astype(bool), image_obj, white_background)
 
-    return Image.fromarray(result_image)
+    import cv2
+
+    edges = cv2.Canny(result_image, 50, 150)
+    # 查找轮廓
+    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # 初始化包围所有外接矩形的大矩形的坐标
+    x_min, y_min, x_max, y_max = float('inf'), float('inf'), -1, -1
+    # 遍历所有外接矩形，更新大矩形的坐标
+    for contour in contours:
+        x, y, w, h = cv2.boundingRect(contour)
+        x_min = min(x_min, x)
+        y_min = min(y_min, y)
+        x_max = max(x_max, x + w)
+        y_max = max(y_max, y + h)
+
+    # 根据大矩形的坐标来裁剪原始图像
+    result_image = image[y_min:y_max, x_min:x_max]
+    # cv2.imshow("", cropped_image)
+    # cv2.waitKey(0)
+
+    return result_image