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"