aidlab/AiDA_Python
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat（新功能）:

Browse Source 
fix（修复bug）:  pose transform mq连接超时bug修复
docs（文档变更）:
refactor（重构）:
test(增加测试):
This commit is contained in:
zchengrong
2025-05-08 17:46:28 +08:00
parent a14e6051b1
commit 293f90f9d3
1 changed files with 31 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
app/service/generate_image/service_pose_transform.py
View File

@@ -29,9 +29,6 @@ logger = logging.getLogger()
class PoseTransformService: class PoseTransformService:
def __init__(self, request_data): def __init__(self, request_data):
if DEBUG is False:
self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
self.channel = self.connection.channel()
self.grpc_client = grpcclient.InferenceServerClient(url=PT_MODEL_URL) self.grpc_client = grpcclient.InferenceServerClient(url=PT_MODEL_URL)
self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True) self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
self.category = "pose_transform" self.category = "pose_transform"
@@ -40,7 +37,8 @@ class PoseTransformService:
self.image = pre_processing_image(request_data.image_url) self.image = pre_processing_image(request_data.image_url)
self.tasks_id = request_data.tasks_id self.tasks_id = request_data.tasks_id
self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:] self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
self.pose_transform_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'gif_url': '', 'video_url': '', 'image_url': ''} self.pose_transform_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'gif_url': '',
'video_url': '', 'image_url': ''}
self.redis_client.set(self.tasks_id, json.dumps(self.pose_transform_data)) self.redis_client.set(self.tasks_id, json.dumps(self.pose_transform_data))
self.redis_client.expire(self.tasks_id, 600) self.redis_client.expire(self.tasks_id, 600)
@@ -55,16 +53,20 @@ class PoseTransformService:
# 第一帧图像 # 第一帧图像
first_image = Image.fromarray(result_data[0]) first_image = Image.fromarray(result_data[0])
first_image_url = upload_first_image(first_image, user_id=self.user_id, category=f"{self.category}_first_img", file_name=f"{self.tasks_id}.png") first_image_url = upload_first_image(first_image, user_id=self.user_id,
category=f"{self.category}_first_img",
file_name=f"{self.tasks_id}.png")
# 上传GIF # 上传GIF
gif_buffer = BytesIO() gif_buffer = BytesIO()
imageio.mimsave(gif_buffer, result_data, format='GIF', fps=5) imageio.mimsave(gif_buffer, result_data, format='GIF', fps=5)
gif_buffer.seek(0) gif_buffer.seek(0)
gif_url = upload_gif(gif_buffer=gif_buffer, user_id=self.user_id, category=f"{self.category}_gif", file_name=f"{self.tasks_id}.gif") gif_url = upload_gif(gif_buffer=gif_buffer, user_id=self.user_id, category=f"{self.category}_gif",
file_name=f"{self.tasks_id}.gif")
# 上传video # 上传video
video_url = upload_video(frames=result_data, user_id=self.user_id, category=f"{self.category}_video", file_name=f"{self.tasks_id}.mp4") video_url = upload_video(frames=result_data, user_id=self.user_id, category=f"{self.category}_video",
file_name=f"{self.tasks_id}.mp4")
self.pose_transform_data['status'] = "SUCCESS" self.pose_transform_data['status'] = "SUCCESS"
self.pose_transform_data['message'] = "success" self.pose_transform_data['message'] = "success"
@@ -82,7 +84,8 @@ class PoseTransformService:
try: try:
pose_num = [self.pose_num] * 1 pose_num = [self.pose_num] * 1
pose_num_obj = np.array(pose_num, dtype="object").reshape((-1, 1)) pose_num_obj = np.array(pose_num, dtype="object").reshape((-1, 1))
input_pose_num = grpcclient.InferInput("pose_num", pose_num_obj.shape, np_to_triton_dtype(pose_num_obj.dtype)) input_pose_num = grpcclient.InferInput("pose_num", pose_num_obj.shape,
np_to_triton_dtype(pose_num_obj.dtype))
input_pose_num.set_data_from_numpy(pose_num_obj) input_pose_num.set_data_from_numpy(pose_num_obj)
image_files = [self.image.astype(np.uint8)] * 1 image_files = [self.image.astype(np.uint8)] * 1
@@ -90,7 +93,8 @@ class PoseTransformService:
input_image_files = grpcclient.InferInput("image_file", image_files_obj.shape, "UINT8") input_image_files = grpcclient.InferInput("image_file", image_files_obj.shape, "UINT8")
input_image_files.set_data_from_numpy(image_files_obj) input_image_files.set_data_from_numpy(image_files_obj)
ctx = self.grpc_client.async_infer(model_name="animatex_1", inputs=[input_pose_num, input_image_files], callback=self.callback, client_timeout=60000) ctx = self.grpc_client.async_infer(model_name="animatex_1", inputs=[input_pose_num, input_image_files],
callback=self.callback, client_timeout=60000)
time_out = 60000 time_out = 60000
while time_out > 0: while time_out > 0:
pose_transform_data, _ = self.read_tasks_status() pose_transform_data, _ = self.read_tasks_status()
@@ -111,9 +115,22 @@ class PoseTransformService:
finally: finally:
dict_pose_transform_data, str_pose_transform_data = self.read_tasks_status() dict_pose_transform_data, str_pose_transform_data = self.read_tasks_status()
if DEBUG is False: if DEBUG is False:
self.channel.basic_publish(exchange='', routing_key=PS_RABBITMQ_QUEUES, body=str_pose_transform_data) publish_status(str_pose_transform_data)
self.connection.close() logger.info(
logger.info(f" [x] Sent to： {PS_RABBITMQ_QUEUES} data：@@@@ {json.dumps(dict_pose_transform_data, indent=4)}") f" [x] Sent to： {PS_RABBITMQ_QUEUES} data：@@@@ {json.dumps(dict_pose_transform_data, indent=4)}")
def publish_status(message):
connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
channel = connection.channel()
channel.queue_declare(queue=PS_RABBITMQ_QUEUES, durable=True)
channel.basic_publish(exchange='',
routing_key=PS_RABBITMQ_QUEUES,
body=json.dumps(message),
properties=pika.BasicProperties(
delivery_mode=2,
))
connection.close()
def infer_cancel(tasks_id): def infer_cancel(tasks_id):
@@ -125,7 +142,8 @@ def infer_cancel(tasks_id):
def pre_processing_image(image_url): 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") image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:],
data_type="PIL")
# 目标图片的尺寸 # 目标图片的尺寸
target_width = 512 target_width = 512
target_height = 768 target_height = 768