diff --git a/app/api/api_pose_transform.py b/app/api/api_pose_transform.py new file mode 100644 index 0000000..fe5fc5a --- /dev/null +++ b/app/api/api_pose_transform.py @@ -0,0 +1,49 @@ +import json +import logging + +from fastapi import APIRouter, BackgroundTasks, HTTPException + +from app.schemas.pose_transform import PoseTransformModel +from app.schemas.response_template import ResponseModel +from app.service.generate_image.service_pose_transform import PoseTransformService, infer_cancel as pose_transform_infer_cancel + +router = APIRouter() +logger = logging.getLogger() + + +@router.post("/pose_transform") +def pose_transform(request_item: PoseTransformModel, background_tasks: BackgroundTasks): + """ + 创建一个具有以下参数的请求体: + - **tasks_id**: 任务id 用于取消生成任务和获取生成结果 + - **image_url**: 被生成图片的S3或minio url地址 + - **pose_id**: 1 + + + 示例参数： + { + "tasks_id": "123-89", + "image_url": "aida-results/result_0000b606-1902-11ef-9424-0242ac180002.png", + "pose_id": "1" + } + """ + try: + logger.info(f"pose_transform request item is : @@@@@@:{json.dumps(request_item.dict())}") + service = PoseTransformService(request_item) + background_tasks.add_task(service.get_result) + except Exception as e: + logger.warning(f"pose_transform Run Exception @@@@@@:{e}") + raise HTTPException(status_code=404, detail=str(e)) + return ResponseModel() + + +@router.get("/pose_transform_cancel/{tasks_id}") +def pose_transform_cancel(tasks_id: str): + try: + logger.info(f"pose_transform_cancel request item is : @@@@@@:{tasks_id}") + data = pose_transform_infer_cancel(tasks_id) + logger.info(f"pose_transform_cancel response @@@@@@:{data}") + except Exception as e: + logger.warning(f"pose_transform_cancel Run Exception @@@@@@:{e}") + raise HTTPException(status_code=404, detail=str(e)) + return ResponseModel(data=data['data']) diff --git a/app/api/api_route.py b/app/api/api_route.py index 45f8567..7fbbc9c 100644 --- a/app/api/api_route.py +++ b/app/api/api_route.py @@ -1,5 +1,6 @@ from fastapi import APIRouter +from app.api import api_agent_generate_image from app.api import api_attribute_retrieve, api_query_image from app.api import api_brand_dna from app.api import api_brighten @@ -9,11 +10,9 @@ from app.api import api_design_pre_processing from app.api import api_generate_image from app.api import api_image2sketch from app.api import api_mannequins_edit +from app.api import api_pose_transform from app.api import api_prompt_generation -from app.api import api_recommendation from app.api import api_super_resolution -from app.api import api_agent_generate_image - from app.api import api_test router = APIRouter() @@ -33,3 +32,4 @@ router.include_router(api_brand_dna.router, tags=['api_brand_dna'], prefix="/api # router.include_router(api_recommendation.router, tags=['api_recommendation'], prefix="/api") router.include_router(api_mannequins_edit.router, tags=['api_mannequins_edit'], prefix="/api") router.include_router(api_agent_generate_image.router, tags=['api_agent_generate_image'], prefix="/api") +router.include_router(api_pose_transform.router, tags=['api_pose_transform'], prefix="/api") diff --git a/app/schemas/pose_transform.py b/app/schemas/pose_transform.py new file mode 100644 index 0000000..045d8b9 --- /dev/null +++ b/app/schemas/pose_transform.py @@ -0,0 +1,7 @@ +from pydantic import BaseModel + + +class PoseTransformModel(BaseModel): + image_url: str + tasks_id: str + pose_id: str diff --git a/app/service/generate_image/service_pose_transform.py b/app/service/generate_image/service_pose_transform.py new file mode 100644 index 0000000..f2948b3 --- /dev/null +++ b/app/service/generate_image/service_pose_transform.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +# -*- coding: UTF-8 -*- +""" +@Project ：trinity_client +@File ：service_pose_transform.py +@Author ：周成融 +@Date ：2023/7/26 12:01:05 +@detail ： +""" +import json +import logging + +import cv2 +import numpy as np +import redis +import tritonclient.grpc as grpcclient +from PIL import Image + +from app.core.config import * +from app.schemas.pose_transform import PoseTransformModel +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 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=GRI_MODEL_URL) + self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True) + self.category = "pose_transform" + self.batch_size = 1 + self.seed = "1" + self.image_url = request_data.image_url + self.image = oss_get_image(bucket=self.image_url.split('/')[0], object_name=self.image_url[self.image_url.find('/') + 1:], data_type="cv2") + 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': 'SUCCESS', 'message': "success", '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: + image = result.as_numpy("generated_inpaint_image") + image_result = Image.fromarray(np.squeeze(image.astype(np.uint8))) + 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) + 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: + image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB) + image = cv2.resize(image, (512, 768)) + images = [image.astype(np.uint8)] * self.batch_size + + image_obj = np.array(images, dtype=np.uint8).reshape((768, 512, 3)) + + input_image = grpcclient.InferInput("input_image", image_obj.shape, "UINT8") + + input_image.set_data_from_numpy(image_obj) + + inputs = [input_image] + # ctx = self.grpc_client.async_infer(model_name=GRI_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", "NO_FACE"]: + # 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=GRI_RABBITMQ_QUEUES, body=str_gen_product_data) + logger.info(f" [x] Sent to： {GRI_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 + + +if __name__ == '__main__': + rd = PoseTransformModel( + tasks_id="123-89", + image_url='aida-results/result_0000b606-1902-11ef-9424-0242ac180002.png', + pose_id="1" + ) + server = PoseTransformService(rd) + print(server.get_result())