feat（新功能）:

fix（修复bug）:
docs（文档变更）:
refactor（重构）:
test(增加测试): 旧版product 测试

app/service/generate_image/service_generate_product_image.py

@@ -1,4 +1,196 @@
-#!/usr/bin/env python
+# #!/usr/bin/env python
+# # -*- coding: UTF-8 -*-
+# """
+# @Project ：trinity_client
+# @File    ：service_att_recognition.py
+# @Author  ：周成融
+# @Date    ：2023/7/26 12:01:05
+# @detail  ：
+# """
+# import json
+# import logging
+# import time
+#
+# import cv2
+# import numpy as np
+# import redis
+# import tritonclient.grpc as grpcclient
+# from PIL import Image
+# from tritonclient.utils import np_to_triton_dtype
+#
+# from app.core.config import *
+# from app.schemas.generate_image import GenerateProductImageModel
+# from app.service.generate_image.utils.upload_sd_image import upload_SDXL_image
+# from app.service.utils.oss_client import oss_get_image
+#
+# logger = logging.getLogger()
+#
+#
+# class GenerateProductImage:
+#     def __init__(self, request_data):
+#         if DEBUG is False:
+#             self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+#             self.channel = self.connection.channel()
+#         # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+#         # self.channel = self.connection.channel()
+#         # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
+#         self.grpc_client = grpcclient.InferenceServerClient(url=GPI_MODEL_URL)
+#         self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+#         self.category = "product_image"
+#         self.image_strength = request_data.image_strength
+#         self.batch_size = 1
+#         self.product_type = request_data.product_type
+#         self.prompt = request_data.prompt
+#         self.image, self.image_size, self.left, self.top = pre_processing_image(request_data.image_url)
+#         self.tasks_id = request_data.tasks_id
+#         self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
+#         self.gen_product_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
+#         self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+#         self.redis_client.expire(self.tasks_id, 600)
+#
+#     def callback(self, result, error):
+#         if error:
+#             self.gen_product_data['status'] = "FAILURE"
+#             self.gen_product_data['message'] = str(error)
+#             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+#         else:
+#             # pil图像转成numpy数组
+#             image = result.as_numpy("generated_inpaint_image")
+#             image_result = Image.fromarray(np.squeeze(image.astype(np.uint8))).resize(self.image_size)
+#             cropped_image = post_processing_image(image_result, self.left, self.top)
+#             image_url = upload_SDXL_image(cropped_image, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
+#             self.gen_product_data['status'] = "SUCCESS"
+#             self.gen_product_data['message'] = "success"
+#             self.gen_product_data['image_url'] = str(image_url)
+#             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+#
+#     def read_tasks_status(self):
+#         status_data = self.redis_client.get(self.tasks_id)
+#         return json.loads(status_data), status_data
+#
+#     def get_result(self):
+#         try:
+#             prompts = [self.prompt] * self.batch_size
+#             self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
+#             self.image = cv2.resize(self.image, (1024, 1024))
+#             images = [self.image.astype(np.uint8)] * self.batch_size
+#
+#             if self.product_type == "single":
+#                 text_obj = np.array(prompts, dtype="object").reshape(-1, 1)
+#                 image_obj = np.array(images, dtype=np.uint8).reshape((-1, 1024, 1024, 3))
+#                 image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape(-1, 1)
+#             else:
+#                 text_obj = np.array(prompts, dtype="object").reshape((1))
+#                 image_obj = np.array(images, dtype=np.uint8).reshape((1024, 1024, 3))
+#                 image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1))
+#
+#             # 假设 prompts、images 和 self.image_strength 已经定义
+#
+#             input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
+#             input_image = grpcclient.InferInput("input_image", image_obj.shape, "UINT8")
+#             input_image_strength = grpcclient.InferInput("image_strength", image_strength_obj.shape, np_to_triton_dtype(image_strength_obj.dtype))
+#
+#             input_text.set_data_from_numpy(text_obj)
+#             input_image.set_data_from_numpy(image_obj)
+#             input_image_strength.set_data_from_numpy(image_strength_obj)
+#
+#             inputs = [input_text, input_image, input_image_strength]
+#
+#             if self.product_type == "single":
+#                 ctx = self.grpc_client.async_infer(model_name="stable_diffusion_xl_cnet_inpaint", inputs=inputs, callback=self.callback)
+#             else:
+#                 ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback)
+#
+#             time_out = 600
+#             while time_out > 0:
+#                 gen_product_data, _ = self.read_tasks_status()
+#                 if gen_product_data['status'] in ["REVOKED", "FAILURE"]:
+#                     ctx.cancel()
+#                     break
+#                 elif gen_product_data['status'] == "SUCCESS":
+#                     break
+#                 time_out -= 1
+#                 time.sleep(0.1)
+#             gen_product_data, _ = self.read_tasks_status()
+#             return gen_product_data
+#         except Exception as e:
+#             self.gen_product_data['status'] = "FAILURE"
+#             self.gen_product_data['message'] = str(e)
+#             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+#             raise Exception(str(e))
+#         finally:
+#             dict_gen_product_data, str_gen_product_data = self.read_tasks_status()
+#             if DEBUG is False:
+#                 self.channel.basic_publish(exchange='', routing_key=GPI_RABBITMQ_QUEUES, body=str_gen_product_data)
+#             logger.info(f" [x] Sent to： {GPI_RABBITMQ_QUEUES} data：@@@@  {json.dumps(dict_gen_product_data, indent=4)}")
+#
+#
+# def infer_cancel(tasks_id):
+#     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+#     data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
+#     gen_product_data = json.dumps(data)
+#     redis_client.set(tasks_id, gen_product_data)
+#     return data
+#
+#
+# def pre_processing_image(image_url):
+#     image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="PIL")
+#     # resize 原图至1024*1024
+#     image = image.resize((int(1024 / image.height * image.width), 1024))
+#
+#     # 原始图片的尺寸
+#     width, height = image.size
+#
+#     new_height, new_width = 1024, 1024
+#     # 创建一个新的画布，大小为添加 padding 后的尺寸，并设置为白色背景
+#     pad_image = Image.new('RGBA', (new_width, new_height), (0, 0, 0, 0))
+#
+#     # 将原始图片粘贴到新的画布中心
+#     left = (new_width - width) // 2
+#     top = (new_height - height) // 2
+#     pad_image.paste(image, (left, top))
+#
+#     # 将画布 resize 成宽度 1024，长度 1024
+#     resized_image = pad_image.resize((1024, 1024))
+#     image_size = (1024, 1024)
+#
+#     if resized_image.mode in ('RGBA', 'LA') or (resized_image.mode == 'P' and 'transparency' in resized_image.info):
+#         # 创建白色背景
+#         background = Image.new("RGB", image_size, (255, 255, 255))
+#         # 将图片粘贴到白色背景上
+#         background.paste(resized_image, mask=resized_image.split()[3])
+#         image = np.array(background)
+#         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+#     return image, image_size, left, top
+#
+#
+# def post_processing_image(image, left, top):
+#     resized_image = image.resize((int(image.width * (768 / image.height)), 768))
+#     # 计算裁剪的坐标
+#     left = (resized_image.width - 512) // 2
+#     upper = 0
+#     right = left + 512
+#     lower = 768
+#
+#     # 进行裁剪
+#     cropped_image = resized_image.crop((left, upper, right, lower))
+#     return cropped_image
+#
+#
+# if __name__ == '__main__':
+#     rd = GenerateProductImageModel(
+#         tasks_id="123-89",
+#         # prompt="",
+#         image_strength=0.7,
+#         prompt="The best quality, masterpiece,outwear, 8K realistic, HUD",
+#         image_url="aida-results/result_53381ada-ac64-11ef-ae9d-0242ac150002.png",
+#         product_type="overall"
+#     )
+#     server = GenerateProductImage(rd)
+#     print(server.get_result())
+
+# 旧版product
+# !/usr/bin/env python
 # -*- coding: UTF-8 -*-
 """
 @Project ：trinity_client
@@ -34,14 +226,14 @@ class GenerateProductImage:
         # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
         # self.channel = self.connection.channel()
         # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
-        self.grpc_client = grpcclient.InferenceServerClient(url=GPI_MODEL_URL)
+        self.grpc_client = grpcclient.InferenceServerClient(url="10.1.1.243:18001")
         self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
         self.category = "product_image"
         self.image_strength = request_data.image_strength
         self.batch_size = 1
         self.product_type = request_data.product_type
         self.prompt = request_data.prompt
-        self.image, self.image_size, self.left, self.top = pre_processing_image(request_data.image_url)
+        self.image = pre_processing_image(request_data.image_url)
         self.tasks_id = request_data.tasks_id
         self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
         self.gen_product_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
@@ -56,9 +248,9 @@ class GenerateProductImage:
         else:
             # pil图像转成numpy数组
             image = result.as_numpy("generated_inpaint_image")
-            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8))).resize(self.image_size)
+            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
-            cropped_image = post_processing_image(image_result, self.left, self.top)
+            # cropped_image = post_processing_image(image_result, self.left, self.top)
-            image_url = upload_SDXL_image(cropped_image, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
+            image_url = upload_SDXL_image(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
             self.gen_product_data['status'] = "SUCCESS"
             self.gen_product_data['message'] = "success"
             self.gen_product_data['image_url'] = str(image_url)
@@ -72,7 +264,7 @@ class GenerateProductImage:
         try:
             prompts = [self.prompt] * self.batch_size
             self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
-            self.image = cv2.resize(self.image, (1024, 1024))
+            # self.image = cv2.resize(self.image, (1024, 1024))
             images = [self.image.astype(np.uint8)] * self.batch_size
 
             if self.product_type == "single":
@@ -81,7 +273,7 @@ class GenerateProductImage:
                 image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape(-1, 1)
             else:
                 text_obj = np.array(prompts, dtype="object").reshape((1))
-                image_obj = np.array(images, dtype=np.uint8).reshape((1024, 1024, 3))
+                image_obj = np.array(images, dtype=np.uint8).reshape((768, 512, 3))
                 image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1))
 
             # 假设 prompts、images 和 self.image_strength 已经定义
@@ -99,7 +291,7 @@ class GenerateProductImage:
             if self.product_type == "single":
                 ctx = self.grpc_client.async_infer(model_name="stable_diffusion_xl_cnet_inpaint", inputs=inputs, callback=self.callback)
             else:
-                ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback)
+                ctx = self.grpc_client.async_infer(model_name="diffusion_ensemble_all", inputs=inputs, callback=self.callback)
 
             time_out = 600
             while time_out > 0:
@@ -135,33 +327,25 @@ def infer_cancel(tasks_id):
 
 def pre_processing_image(image_url):
     image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="PIL")
-    # resize 原图至1024*1024
+    # 调整图片高度为768像素，保持宽高比
-    image = image.resize((int(1024 / image.height * image.width), 1024))
-
-    # 原始图片的尺寸
     width, height = image.size
+    new_height = 768
+    new_width = int(width * (new_height / height))
+    resized_image = image.resize((new_width, new_height))
 
-    new_height, new_width = 1024, 1024
+    # 创建一个512x768的透明图片
-    # 创建一个新的画布，大小为添加 padding 后的尺寸，并设置为白色背景
+    result_image = Image.new("RGBA", (512, 768), (0, 0, 0, 0))
-    pad_image = Image.new('RGBA', (new_width, new_height), (0, 0, 0, 0))
 
-    # 将原始图片粘贴到新的画布中心
+    # 计算需要粘贴的位置，使图片居中
-    left = (new_width - width) // 2
+    x_offset = (512 - new_width) // 2
-    top = (new_height - height) // 2
+    y_offset = 0
-    pad_image.paste(image, (left, top))
 
-    # 将画布 resize 成宽度 1024，长度 1024
+    # 将调整大小后的图片粘贴到透明图片上
-    resized_image = pad_image.resize((1024, 1024))
+    result_image.paste(resized_image, (x_offset, y_offset))
-    image_size = (1024, 1024)
 
-    if resized_image.mode in ('RGBA', 'LA') or (resized_image.mode == 'P' and 'transparency' in resized_image.info):
+    image = np.array(result_image)
-        # 创建白色背景
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        background = Image.new("RGB", image_size, (255, 255, 255))
+    return image
-        # 将图片粘贴到白色背景上
-        background.paste(resized_image, mask=resized_image.split()[3])
-        image = np.array(background)
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    return image, image_size, left, top
 
 
 def post_processing_image(image, left, top):