import gradio as gr
import os


os.system("git clone https://github.com/megvii-research/NAFNet")
os.system("mv NAFNet/* ./")
os.system("mv *.pth experiments/pretrained_models/")
os.system("python3 setup.py develop --no_cuda_ext --user")


def inference(image, task):
    if not os.path.exists('tmp'):
      os.system('mkdir tmp')
    image.save("tmp/lq_image.png", "PNG")
    
    if task == 'Denoising':
      os.system("python basicsr/demo.py -opt options/test/SIDD/NAFNet-width64.yml --input_path ./tmp/lq_image.png --output_path ./tmp/image.png")
  
    if task == 'Deblurring':
      os.system("python basicsr/demo.py -opt options/test/REDS/NAFNet-width64.yml --input_path ./tmp/lq_image.png --output_path ./tmp/image.png")
  
    return 'tmp/image.png'
   
title = "DCGAN"
description = "DCGAN 的 Gradio 演示：用于图像恢复的非线性无激活网络。DCGAN 在三个任务上实现了最先进的性能：图像去噪。在这里，提供了一个图像去噪的演示。要使用它，只需上传您的图像，或单击其中一个示例以加载它们。由于此演示使用 CPU，因此推理需要一些时间。"
#article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2204.04676' target='_blank'>Simple Baselines for Image Restoration</a> | <a href='https://arxiv.org/abs/2204.08714' target='_blank'>NAFSSR: Stereo Image Super-Resolution Using NAFNet</a>  | <a href='https://github.com/megvii-research/NAFNet' target='_blank'> Github Repo</a></p>"


examples = [['demo/noisy.png', 'Denoising'],
            ['demo/blurry.jpg', 'Deblurring']]
            
iface = gr.Interface(
    inference, 
    [gr.inputs.Image(type="pil", label="Input"),
    gr.inputs.Radio(["Denoising", "Deblurring"], default="Denoising", label='task'),], 
    gr.outputs.Image(type="file", label="Output"),
    title=title,
    description=description,
    # article=article,
    enable_queue=True,
    examples=examples
    )
iface.launch(debug=True,enable_queue=True)