feat generate sketch 新增 服装类别识别

app/service/generate_image/service.py

@@ -52,6 +52,7 @@ class GenerateImage:
         self.batch_size = 1
         self.category = request_data.category
         self.index = 0
+        self.gender = request_data.gender
         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)
@@ -84,11 +85,11 @@ class GenerateImage:
                 # 污点检测
                 is_smudge, not_smudge_image = stain_detection(remove_bg_image)
                 # 类型识别
-                category, scores, not_smudge_image = generate_category_recognition(image_result)
+                category, scores, not_smudge_image = generate_category_recognition(image=image_result, gender=self.gender)
                 self.generate_data['category'] = str(category)
                 image_result = not_smudge_image
             if is_smudge:  # 无污点
-                image_result = adjust_contrast(image_result)
+                # image_result = adjust_contrast(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")
                 # logger.info(f"upload image SUCCESS ： {image_url}")
                 self.generate_data['status'] = "SUCCESS"