AiDA_Python/app/service/generate_image/service_generate_image.py

#!/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 minio
import numpy as np
import redis
import tritonclient.grpc as grpcclient
from tritonclient.utils import np_to_triton_dtype

from app.core.config import *
from app.schemas.generate_image import GenerateImageModel
from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust
from app.service.generate_image.utils.mq import publish_status
from app.service.generate_image.utils.upload_sd_image import upload_png_sd
from app.service.utils.oss_client import oss_get_image

logger = logging.getLogger()


class GenerateImage:
    def __init__(self, request_data):
        self.version = request_data.version
        if request_data.version == "fast":
            self.grpc_client = grpcclient.InferenceServerClient(url=FAST_GI_MODEL_URL)
        else:
            self.grpc_client = grpcclient.InferenceServerClient(url=GI_MODEL_URL)

        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
        if request_data.mode == "img2img":
            # cv2 读图片是BGR PIL读图片是RGB
            self.image = self.get_image(request_data.image_url)
        else:
            self.image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)

        self.prompt = request_data.prompt
        self.tasks_id = request_data.tasks_id
        self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
        self.mode = request_data.mode
        self.batch_size = 1
        self.category = request_data.category
        if self.category == "sketch":
            self.prompt = f"{self.category},{self.prompt}"
        self.index = 0
        self.gender = request_data.gender
        self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': '', 'category': ''}
        self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
        self.redis_client.expire(self.tasks_id, 600)

    def get_image(self, image_url):
        # Get data of an object.
        # Read data from response.
        # read image use cv2
        try:
            # response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
            # image_file = BytesIO(response.data)
            # image_array = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
            # image_cv2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
            # image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)

            image_cv2 = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="cv2")
            image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image_rbg, (1024, 1024))
        except minio.error.S3Error:
            image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
        return image

    def callback(self, result, error):
        if error:
            self.generate_data['status'] = "FAILURE"
            self.generate_data['message'] = str(error)
            # self.generate_data['data'] = str(error)
            self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
        else:
            # pil图像转成numpy数组
            image = result.as_numpy("generated_image")
            image_result = cv2.cvtColor(np.squeeze(image.astype(np.uint8)), cv2.COLOR_RGB2BGR)
            is_smudge = True
            if self.category == "sketch":
                if self.version == "fast":
                    # 色阶调整
                    cutoff = 1
                    levels_img = autoLevels(image_result, cutoff)
                    # 亮度调整
                    luminance = luminance_adjust(0.3, levels_img)
                    # 去背景
                    remove_bg_image = remove_background(luminance)
                    # 人脸检测
                    # if face_detect_pic(remove_bg_image, self.user_id, self.category, self.tasks_id) > 0:
                    #     is_smudge = False
                    # else:
                    # 污点/
                    is_smudge, not_smudge_image = stain_detection(remove_bg_image, self.user_id, self.category, self.tasks_id)
                    # 类型识别
                    category, scores, not_smudge_image = generate_category_recognition(image=remove_bg_image, gender=self.gender)
                    self.generate_data['category'] = str(category)
                    image_result = not_smudge_image
                else:
                    category, scores, not_smudge_image = generate_category_recognition(image=image_result, gender=self.gender)
                    self.generate_data['category'] = str(category)
                    image_result = not_smudge_image
            if is_smudge:  # 无污点
                # image_result = adjust_contrast(image_result)
                image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
                # logger.info(f"upload image SUCCESS ： {image_url}")
                self.generate_data['status'] = "SUCCESS"
                self.generate_data['message'] = "success"
                self.generate_data['image_url'] = str(image_url)
                self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
            else:  # 有污点 保存图片到本地 测试用
                self.generate_data['status'] = "SUCCESS"
                self.generate_data['message'] = "success"
                self.generate_data['image_url'] = str(GI_SYS_IMAGE_URL)
                self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
            # logger.info(f"stain_detection result : {self.generate_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
            modes = [self.mode] * self.batch_size
            images = [self.image.astype(np.float16)] * self.batch_size

            text_obj = np.array(prompts, dtype="object").reshape((-1, 1))
            mode_obj = np.array(modes, dtype="object").reshape((-1, 1))
            image_obj = np.array(images, dtype=np.float16).reshape((-1, 1024, 1024, 3))

            input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
            input_image = grpcclient.InferInput("input_image", image_obj.shape, np_to_triton_dtype(image_obj.dtype))
            input_mode = grpcclient.InferInput("mode", mode_obj.shape, np_to_triton_dtype(mode_obj.dtype))

            input_text.set_data_from_numpy(text_obj)
            input_image.set_data_from_numpy(image_obj)
            input_mode.set_data_from_numpy(mode_obj)

            inputs = [input_text, input_image, input_mode]
            if self.version == "fast":
                ctx = self.grpc_client.async_infer(model_name=FAST_GI_MODEL_NAME, inputs=inputs, callback=self.callback, priority=1)
            else:
                ctx = self.grpc_client.async_infer(model_name=GI_MODEL_NAME, inputs=inputs, callback=self.callback, priority=1)

            time_out = 600
            generate_data = None
            while time_out > 0:
                generate_data, _ = self.read_tasks_status()
                if generate_data['status'] in ["REVOKED", "FAILURE"]:
                    ctx.cancel()
                    break
                elif generate_data['status'] == "SUCCESS":
                    break
                time_out -= 1
                time.sleep(0.1)
            return generate_data
        except Exception as e:
            self.generate_data['status'] = "FAILURE"
            self.generate_data['message'] = str(e)
            self.redis_client.set(self.tasks_id, json.dumps(self.generate_data))
            raise Exception(str(e))
        finally:
            dict_generate_data, str_generate_data = self.read_tasks_status()
            if not DEBUG:
                publish_status(str_generate_data, GI_RABBITMQ_QUEUES)


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'}
    generate_data = json.dumps(data)
    redis_client.set(tasks_id, generate_data)
    return data


if __name__ == '__main__':
    rd = GenerateImageModel(
        tasks_id="123-89",
        prompt="Women's clothing ,dress,technical drawing style, clean line art, no shading, no texture, flat sketch, no human body, no face, centered composition, pure white background, single garmentsingle garment only, front flat view",
        image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg",
        mode='txt2img',
        category="test",
        gender="male",
        version="high"
    )
    server = GenerateImage(rd)
    print(server.get_result())