import os
from models import create_model
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(110)



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 = "애니메이션풍 이미지를 볼짤콘 스타일로 바꾸는 Pix2pix모델입니다. 이미지 업로드 후 나오는 우측 상단에 연필모양 클릭해서 얼굴 쪽을 정사각형으로 자르고 submit 하면 됩니다. 큰 이미지는 화질 개선 체크 하시고 작은 이미지는 자동으로 개선되니 상관 X. 볼짤콘이 얼굴밖에 없다보니 얼굴이 중앙에 있고 안면이 크게 나올수록 잘 되는 느낌입니다. 문의 krkmfn@gmail.com "

examples=[['momoi_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()