aidlab/AiDA_Python
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat 新增 生成sketch时对图片清理背景，剔除带有污点的结果图

Browse Source
This commit is contained in:
zchen
2024-04-23 15:53:24 +08:00
parent 2c2b0d3b0d
commit e5adb08d04
1 changed files with 17 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
app/service/generate_image/service.py
View File

@@ -75,26 +75,28 @@ class GenerateImage:
self.generate_data['data'] = str(error) self.generate_data['data'] = str(error)
self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
else: else:
logger.info(1111111111111111111111111111)
image_result = result.as_numpy("generated_image")[0] image_result = result.as_numpy("generated_image")[0]
is_smudge = True
if self.category == "sketch": if self.category == "sketch":
# 去背景 # 去背景
remove_bg_image = remove_background(np.asarray(image_result)) remove_bg_image = remove_background(np.asarray(image_result))
# 污点检测 # 污点检测
is_smudge, not_smudge_image = stain_detection(remove_bg_image) is_smudge, not_smudge_image = stain_detection(remove_bg_image)
if is_smudge: # 无污点 image_result = not_smudge_image
image_result = not_smudge_image if is_smudge: # 无污点
image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png") image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")
logger.info(f"upload image SUCCESS ： {image_url}") logger.info(f"upload image SUCCESS ： {image_url}")
self.generate_data['status'] = "SUCCESS" self.generate_data['status'] = "SUCCESS"
self.generate_data['message'] = "success" self.generate_data['message'] = "success"
self.generate_data['data'] = str(image_url) self.generate_data['data'] = str(image_url)
self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
else: # 有污点 else: # 有污点
self.generate_data['status'] = "SUCCESS" self.generate_data['status'] = "SUCCESS"
self.generate_data['message'] = "success" self.generate_data['message'] = "success"
self.generate_data['data'] = str(GI_SYS_IMAGE_URL) self.generate_data['data'] = str(GI_SYS_IMAGE_URL)
self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
logger.info(f"stain_detection result : {self.generate_data}") logger.info(f"stain_detection result : {self.generate_data}")
def read_tasks_status(self): def read_tasks_status(self):
status_data = self.redis_client.get(self.tasks_id) status_data = self.redis_client.get(self.tasks_id)
@@ -131,6 +133,7 @@ class GenerateImage:
generate_data = None generate_data = None
while time_out > 0: while time_out > 0:
generate_data, _ = self.read_tasks_status() generate_data, _ = self.read_tasks_status()
logger.info(generate_data)
if generate_data['status'] in ["REVOKED", "FAILURE"]: if generate_data['status'] in ["REVOKED", "FAILURE"]:
ctx.cancel() ctx.cancel()
break break