# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import os
os.system('pip install -r requirements.txt')

import cv2
import gradio as gr

from ctrnet_infer import CTRNetInfer


def inference(img_path):
    img = cv2.imread(img_path)
    pred = ctrnet(img)
    pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
    return pred


model_path = 'models/CTRNet_G.onnx'
ctrnet = CTRNetInfer(model_path)

title = 'CTRNet Demo'
description = '''This is the demo for the paper “Don't Forget Me: Accurate Background Recovery for Text Removal via Modeling Local-Global Context”. Github Repo: https://github.com/lcy0604/CTRNet'''
css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
examples = [['images/1.jpg'], ['images/2.jpg'], ['images/4.jpg']]

gr.Interface(
    inference,
    inputs=[
        gr.inputs.Image(type='filepath', label='Input'),
    ],
    outputs=[
        gr.outputs.Image(type='filepath', label='Output_image'),
    ],
    title=title,
    description=description,
    examples=examples,
    css=css,
    allow_flagging='never',
    enable_queue=True
    ).launch(debug=True, enable_queue=True)