feat（新功能）:
fix（修复bug）: docs（文档变更）: refactor（重构）: test(增加测试): 旧版product 测试
This commit is contained in:
zhouchengrong
@@ -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)
|
||||
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_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
|
||||
# cropped_image = post_processing_image(image_result, self.left, self.top)
|
||||
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
|
||||
image = image.resize((int(1024 / image.height * image.width), 1024))
|
||||
|
||||
# 原始图片的尺寸
|
||||
# 调整图片高度为768像素，保持宽高比
|
||||
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
|
||||
# 创建一个新的画布，大小为添加 padding 后的尺寸，并设置为白色背景
|
||||
pad_image = Image.new('RGBA', (new_width, new_height), (0, 0, 0, 0))
|
||||
# 创建一个512x768的透明图片
|
||||
result_image = Image.new("RGBA", (512, 768), (0, 0, 0, 0))
|
||||
|
||||
# 将原始图片粘贴到新的画布中心
|
||||
left = (new_width - width) // 2
|
||||
top = (new_height - height) // 2
|
||||
pad_image.paste(image, (left, top))
|
||||
# 计算需要粘贴的位置，使图片居中
|
||||
x_offset = (512 - new_width) // 2
|
||||
y_offset = 0
|
||||
|
||||
# 将画布 resize 成宽度 1024，长度 1024
|
||||
resized_image = pad_image.resize((1024, 1024))
|
||||
image_size = (1024, 1024)
|
||||
# 将调整大小后的图片粘贴到透明图片上
|
||||
result_image.paste(resized_image, (x_offset, y_offset))
|
||||
|
||||
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
|
||||
image = np.array(result_image)
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
return image
|
||||
|
||||
|
||||
def post_processing_image(image, left, top):
|
||||
|
||||
Reference in New Issue
Block a user