#pip install "modelscope[cv]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
#pip install gradio
#pip install tensorflow

from tqdm import tqdm
# from skimage import io
import datetime
import os
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from translate import Translator
from gradio_client import Client
import json

# 初始化Translator对象,指定源语言和目标语言
translator = Translator(from_lang="zh", to_lang="en")

#获取当前北京时间
utc_dt = datetime.datetime.utcnow()
beijing_dt = utc_dt.astimezone(datetime.timezone(datetime.timedelta(hours=8)))
formatted = beijing_dt.strftime("%Y-%m-%d_%H")
print(f"北京时间: {beijing_dt.year}年{beijing_dt.month}月{beijing_dt.day}日 "
      f"{beijing_dt.hour}时{beijing_dt.minute}分{beijing_dt.second}秒")
#创建作品存放目录
works_path = 'works_text_image_api/' + formatted
if not os.path.exists(works_path):
  os.makedirs(works_path)
print('作品目录：' + works_path)
#创建用户上传图片存放目录
user_upload_path = 'user_upload/' + formatted
if not os.path.exists(user_upload_path):
  os.makedirs(user_upload_path)
print('用户图片目录：' + user_upload_path)

def get_size(h, w, max = 720):
    if min(h, w) > max:
        if h > w:
            h, w = int(max * h / w), max
        else:
            h, w = max, int(max * w / h)
    
    return h, w
        
def inference(original_prompt: str, image: Image) -> Image:
  #调整图片尺寸，避免过大导致处理耗时过久
  w, h = image.size
  print(f'原图片宽：{w}，高：{h}')
  h, w = get_size(h, w, 720)
  image = image.resize((w, h))
  print(f'调整尺寸后图片宽：{w}，高：{h}')

  print('图片描述：' + original_prompt)
  translate_prompt = translator.translate(original_prompt) #翻译为英文
  print('translate_prompt：' + translate_prompt)
    
  utc_dt = datetime.datetime.utcnow()
  beijing_dt = utc_dt.astimezone(datetime.timezone(datetime.timedelta(hours=8)))
  formatted = beijing_dt.strftime("%Y-%m-%d_%H-%M-%S.%f")
  image_path = user_upload_path + '/' + formatted + '.png'
  print('用户图片：' + image_path)
  image.save(image_path)
  
  # https://huggingface.co/spaces/hysts/ControlNet-v1-1
  client = Client("https://hysts-controlnet-v1-1.hf.space/")
  result = client.predict(
        image_path,    # str (filepath or URL to image) in 'parameter_98' Image component
        translate_prompt,    # str  in 'Prompt' Textbox component
        "masterpiece, best quality, extremely detailed",    # str  in 'Additional prompt' Textbox component
        "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",    # str  in 'Negative prompt' Textbox component
        1,    # int | float (numeric value between 1 and 1) in 'Number of images' Slider component
        512,    # int | float (numeric value between 256 and 512) in 'Image resolution' Slider component
        512,    # int | float (numeric value between 128 and 512) in 'Preprocess resolution' Slider component
        20,    # int | float (numeric value between 1 and 100) in 'Number of steps' Slider component
        9.0,    # int | float (numeric value between 0.1 and 30.0) in 'Guidance scale' Slider component
        706138,    # int | float (numeric value between 0 and 1000000) in 'Seed' Slider component
        "HED",    # str  in 'Preprocessor' Radio component
        api_name="/softedge"
  )

  print(result)
  result += '/captions.json'

  with open(result) as f:
    data = json.load(f)

  # data 为Python对象
  print(data)
  
  i = 0
  for key in data.keys():
    if i == 1:
        result_path = key
        break
    
    i += 1

  res_img = Image.open(result_path)
  print('作品：' + result_path)

  # 加载字体，设置字体大小
  font_path = 'ttf/WawaSC-Regular.otf'
  font_size = 50
  font = ImageFont.truetype(font_path, font_size)
  text = 'by 宁侠'

  x0, y0, x1, y1 = font.getbbox(text)
  text_width = x1 - x0
  text_height = (y1 - y0)*2

  watermark = Image.new('RGBA', (text_width, text_height))
  draw = ImageDraw.Draw(watermark)
  draw.text((0,0), text, font=font, fill=(255,255,255)) #阿根廷蓝：112,171,221

  w, h = res_img.size
  res_img.paste(watermark, (w - text_width - 10, h - text_height), watermark)

  return res_img


css_style = "#fixed_size_img {height: 240px;} "

title = "人像创作 by宁侠"
description = '''
我们提供的服务能够快速高效地将您提供的人像图片转化为栩栩如生的肖像图，您只需简单地输入图片描述，我们的服务便会根据您的要求对图片进行处理，让您获得一张高质量的肖像图。我们期待着为您提供最好的服务，并让您的体验更加愉快。
'''
examples_path = 'examples/'
examples = [[examples_path + 'input1.png'], [examples_path + 'input2.png'], [examples_path + 'input3.png'], [examples_path + 'input4.png']]

with gr.Blocks(title=title, css=css_style) as demo:
    gr.HTML('''
      <div style="text-align: center; max-width: 720px; margin: 0 auto;">
                  <div
                    style="
                      display: inline-flex;
                      align-items: center;
                      gap: 0.8rem;
                      font-size: 1.75rem;
                    "
                  >
                    <h1 style="font-family: PingFangSC; font-weight: 500; line-height: 1.5em; font-size: 32px; margin-bottom: 7px;">
                      人像创作
                    </h1>
                    <h1 style="font-family: PingFangSC; font-weight: 500; line-height: 1.5em; font-size: 16px; margin-bottom: 7px;">
                      by宁侠
                    </h1>
                  </div>
                </div>
      ''')

    gr.Markdown(description)
    with gr.Row():
        original_prompt = gr.Textbox(label="请输入图片描述", value="英俊青年的油画，杰作")
    with gr.Row():
        img_input = gr.Image(label="图片", type="pil", elem_id="fixed_size_img")
        img_output = gr.Image(label="作品", type="pil", elem_id="fixed_size_img")
    with gr.Row():
        btn_submit = gr.Button(value="一键创作", elem_id="blue_btn")
        # btn_clear = gr.Button(value="清除")

    examples = gr.Examples(examples=examples, inputs=[img_input], outputs=img_output)
    btn_submit.click(inference, inputs=[original_prompt, img_input], outputs=img_output)
    # btn_clear清除画布

demo.queue(api_open=False).launch(debug=True)