diff --git a/app/service/generate_image/service.py b/app/service/generate_image/service.py index b66c7a5..ea39a21 100644 --- a/app/service/generate_image/service.py +++ b/app/service/generate_image/service.py @@ -81,18 +81,19 @@ class GenerateImage: remove_bg_image = remove_background(np.asarray(image_result)) # 污点检测 is_smudge, not_smudge_image = stain_detection(remove_bg_image) - if is_smudge is False: - self.generate_data['status'] = "SUCCESS" - self.generate_data['message'] = "success" - self.generate_data['data'] = GI_SYS_IMAGE_URL - self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) - else: + if is_smudge: # 无污点 image_result = not_smudge_image 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" self.generate_data['data'] = str(image_url) self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) + else: # 有污点 + self.generate_data['status'] = "SUCCESS" + self.generate_data['message'] = "success" + self.generate_data['data'] = str(GI_SYS_IMAGE_URL) + self.redis_client.set(self.tasks_id, json.dumps(self.generate_data)) + logger.info(f"stain_detection result : {self.generate_data}") def read_tasks_status(self): status_data = self.redis_client.get(self.tasks_id)