feat generate sketch 新增 服装类别识别

This commit is contained in:
zhouchengrong
2024-04-24 11:24:38 +08:00
parent 7f46294c6a
commit 679c854b01

View File

@@ -52,7 +52,7 @@ class GenerateImage:
self.batch_size = 1
self.category = request_data.category
self.index = 0
self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'data': {"image_url": "", "category": ""}}
self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': '', 'category': ''}
self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
self.redis_client.expire(self.tasks_id, 600)
@@ -85,7 +85,7 @@ class GenerateImage:
is_smudge, not_smudge_image = stain_detection(remove_bg_image)
# 类型识别
category, scores, not_smudge_image = generate_category_recognition(image_result)
self.generate_data['data']['category'] = str(category)
self.generate_data['category'] = str(category)
image_result = not_smudge_image
if is_smudge: # 无污点
image_result = adjust_contrast(image_result)
@@ -93,12 +93,12 @@ class GenerateImage:
# logger.info(f"upload image SUCCESS {image_url}")
self.generate_data['status'] = "SUCCESS"
self.generate_data['message'] = "success"
self.generate_data['data']['image_url'] = str(image_url)
self.generate_data['image_url'] = 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']['image_url'] = str(GI_SYS_IMAGE_URL)
self.generate_data['image_url'] = 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}")