import gradio as gr
import os,re
import numpy as np
import requests
import cv2
import base64
import json
from io import BytesIO
from PIL import Image

# api_service_url = 'http://127.0.0.1:1234'
# api_service_url = 'http://i-2.gpushare.com:53412'
api_service_url = 'http://service.aicloud.fit:7393/api/hairStyle/v1'

class ChangeHairGui():
    # def __init__(self):
    #     a = 1
    def change_hair(self, user_img, hair_img):

        if user_img is None or user_img is None:
            return None
        #
        # if user_img.shape != hair_img.shape:
        #     hair_img = cv2.resize(hair_img, (user_img.shape[1], user_img.shape[0]))
        #
        #     alpha = 0.5  # 图像1的权重
        #     beta = 0.5  # 图像2的权重
        #     gamma = 0  # 亮度调整常量（通常为0）
        #
        #     result_image = cv2.addWeighted(user_img, alpha, hair_img, beta, gamma)


        # 将 Pillow 图像转换为 OpenCV 格式（BGR）
        user_img = cv2.cvtColor(np.array(user_img), cv2.COLOR_RGB2BGR)
        hair_img = cv2.cvtColor(np.array(hair_img), cv2.COLOR_RGB2BGR)

        # 将模板图片转换为base64格式
        retval, user_bytes = cv2.imencode('.jpg', user_img)
        encoded_user_image = base64.b64encode(user_bytes).decode('utf-8')
        retval, hair_bytes = cv2.imencode('.jpg', hair_img)
        encoded_hair_image = base64.b64encode(hair_bytes).decode('utf-8')

        url = api_service_url

        # 请求换发型接口
        payload = json.dumps({
            "user_img_base64": encoded_user_image,
            "hair_ref_img_base64": encoded_hair_image
        })
        headers = {
            'Content-Type': 'application/json'
        }
        print('请求api_service.py发送请求！！')
        response = requests.request("POST", url, headers=headers, data=payload)
        print('请求api_service.py发送请求成功！！')
        if response.status_code != 200:
            raise RuntimeError(f"Failed to send request to API service. Status code: {response.status_code}")

        ret_image_b64 = response.json().get('result')

        if ret_image_b64 is None:
            raise RuntimeError(f"ret image failed!")

        image_array = np.frombuffer(base64.b64decode(ret_image_b64), np.uint8)
        result_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
        result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)

        dst_size = max(result_image.shape[0], result_image.shape[1])
        M = cv2.getRotationMatrix2D((result_image.shape[1] / 2, result_image.shape[0] / 2), 0, 1)
        M[:, 2] += np.float32([dst_size / 2 - result_image.shape[1] / 2, dst_size / 2 - result_image.shape[0] / 2])
        result_image = cv2.warpAffine(result_image, M, (dst_size, dst_size), borderValue=(255, 255, 255))

        return result_image



    def start_gui(self):
        with gr.Blocks() as demo:
            with gr.Row():
                gr.Markdown("# 数字力场换发型效果展示")

            # 换脸
            with gr.Tab("换发型"):
                with gr.Row():
                    gr.Markdown("#换发型")
                with gr.Row():
                    with gr.Column():
                        user_img = gr.Image(label="请上传用户图片", type="numpy", height=384, width=384)
                    with gr.Column():
                        hair_img = gr.Image(label="请上传发型图片", type="numpy", height=384, width=384)
                    with gr.Column():
                        output_img = gr.Image(label="结果展示", type="numpy", height=384, width=384, format='png')
                with gr.Column():
                    run_button = gr.Button(value="提交")
                    run_button.click(self.change_hair, inputs=[user_img, hair_img], outputs=[output_img])

        demo.launch(server_name='0.0.0.0', server_port=8080)

if __name__ == '__main__':

    demo = ChangeHairGui()
    demo.start_gui()