import os
from PIL import Image
from imageio import imwrite
import gradio as gr
from torchvision import transforms
import easydict
import torch
import numpy as np
import cv2
from subprocess import call


def run_cmd(command):
    try:
        print(command)
        call(command, shell=True)
    except KeyboardInterrupt:
        print("Process interrupted")
        sys.exit(1)
run_cmd("pip install basicsr")
run_cmd("pip freeze")

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
import train



esrgan_model = train.build_esrgan()
pix2pix_model = train.build_pix2pix()



def pipeline(img, vivid):
  img = train.image_preprosses(img, vivid)
  if type(img) is np.ndarray:
    img, _ = esrgan_model.enhance(img, outscale=4)
    img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
  img = train.test_pix2pix(img, pix2pix_model)
  imwrite('filename.png', img)
  
  return 'filename.png'


description = "애니메이션풍 이미지를 볼짤콘 스타일로 바꾸는 Pix2pixHD모델입니다. 이미지 업로드 후 나오는 우측 상단에 연필모양 클릭해서 얼굴 쪽을 정사각형으로 자르고 submit 하면 됩니다. 출력 이미지가 흐리게 나오면 화질개선 체크하시면 됩니다."


examples=[['siroko_out.png',False]]
gr.Interface(pipeline, [gr.inputs.Image(type="pil"),gr.inputs.Checkbox(default=False, label="화질 개선")], gr.outputs.Image(type="file"),description=description,allow_flagging=False,examples=examples,allow_screenshot=False,enable_queue=False).launch()