'''
Author: Egrt
Date: 2022-01-13 13:34:10
LastEditors: Egrt
LastEditTime: 2022-01-13 13:48:57
FilePath: \LicenseGAN\app.py
'''
import os
os.system('pip install pytorch')
os.system('pip install gradio==2.5.3')
from PIL import Image
from esrgan import ESRGAN
import gradio as gr

esrgan = ESRGAN()

# --------模型推理---------- #
def inference(img):
    lr_shape = [12, 24]
    img = img.resize((lr_shape[1], lr_shape[0]), Image.BICUBIC)
    r_image = esrgan.generate_1x1_image(img)
    return r_image

# --------网页信息---------- #  
title = "车牌超分辨率重建"
description = "使用生成对抗网络对低分辨率车牌图片进行八倍的超分辨率重建，能够有效的恢复出车牌号。  @西南科技大学智能控制与图像处理研究室"
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.10257' target='_blank'>LicenseGAN: Image Restoration Using Swin Transformer</a> | <a href='https://github.com/JingyunLiang/SwinIR' target='_blank'>Github Repo</a></p>"
example_img_dir  = 'img'
example_img_name = os.listdir(example_img_dir)
examples=[[os.path.join(example_img_dir, image_path)] for image_path in example_img_name if image_path.endswith('.jpg')]
gr.Interface(
    inference, 
    [gr.inputs.Image(type="pil", label="Input")], 
    gr.outputs.Image(type="pil", label="Output"),
    title=title,
    description=description,
    article=article,
    enable_queue=True,
    examples=examples
    ).launch(debug=True)