From 293f90f9d387cf382f8c6b3def369c8c938513f4 Mon Sep 17 00:00:00 2001 From: zchengrong <124802516+zchengrong@users.noreply.github.com> Date: Thu, 8 May 2025 17:46:28 +0800 Subject: [PATCH] =?UTF-8?q?feat=EF=BC=88=E6=96=B0=E5=8A=9F=E8=83=BD?= =?UTF-8?q?=EF=BC=89:=20fix=EF=BC=88=E4=BF=AE=E5=A4=8Dbug=EF=BC=89:=20=20p?= =?UTF-8?q?ose=20transform=20mq=E8=BF=9E=E6=8E=A5=E8=B6=85=E6=97=B6bug?= =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20docs=EF=BC=88=E6=96=87=E6=A1=A3=E5=8F=98?= =?UTF-8?q?=E6=9B=B4=EF=BC=89:=20refactor=EF=BC=88=E9=87=8D=E6=9E=84?= =?UTF-8?q?=EF=BC=89:=20test(=E5=A2=9E=E5=8A=A0=E6=B5=8B=E8=AF=95):?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../generate_image/service_pose_transform.py | 44 +++++++++++++------ 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/app/service/generate_image/service_pose_transform.py b/app/service/generate_image/service_pose_transform.py index 07da8de..2bd81ac 100644 --- a/app/service/generate_image/service_pose_transform.py +++ b/app/service/generate_image/service_pose_transform.py @@ -29,9 +29,6 @@ logger = logging.getLogger() class PoseTransformService: 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.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True) self.category = "pose_transform" @@ -40,7 +37,8 @@ class PoseTransformService: 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.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.expire(self.tasks_id, 600) @@ -55,16 +53,20 @@ class PoseTransformService: # 第一帧图像 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_buffer = BytesIO() imageio.mimsave(gif_buffer, result_data, format='GIF', fps=5) 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_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['message'] = "success" @@ -82,7 +84,8 @@ class PoseTransformService: try: pose_num = [self.pose_num] * 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) 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.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 while time_out > 0: pose_transform_data, _ = self.read_tasks_status() @@ -111,9 +115,22 @@ class PoseTransformService: finally: dict_pose_transform_data, str_pose_transform_data = self.read_tasks_status() if DEBUG is False: - self.channel.basic_publish(exchange='', routing_key=PS_RABBITMQ_QUEUES, body=str_pose_transform_data) - self.connection.close() - logger.info(f" [x] Sent to： {PS_RABBITMQ_QUEUES} data：@@@@ {json.dumps(dict_pose_transform_data, indent=4)}") + publish_status(str_pose_transform_data) + logger.info( + 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): @@ -125,7 +142,8 @@ 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") + image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], + data_type="PIL") # 目标图片的尺寸 target_width = 512 target_height = 768