import os
os.system('pip install --upgrade --no-cache-dir gdown')
os.system('gdown -O ./output/ctw/model_ctw.pth 1ydYjuxJglLAJRVWjj6sKXM3aBL2pXRpG')
os.system('gdown -O ./workdir.zip 1mYM_26qHUom_5NU7iutHneB_KHlLjL5y')
os.system('unzip workdir.zip')
os.system('pip install "git+https://github.com/philferriere/cocoapi.git#egg=pycocotools&subdirectory=PythonAPI"')
os.system('python setup.py build develop --user')

import cv2
import pandas as pd
import gradio as gr

from det_demo import DetDemo
from maskrcnn_benchmark.config import cfg

from demo import get_model, preprocess, postprocess, load
from utils import Config, Logger, CharsetMapper
import torch


def infer(filepath):
    config = Config('configs/rec/train_abinet.yaml')
    config.model_vision_checkpoint = None
    model = get_model(config)
    model = load(model, 'workdir/train-abinet/best-train-abinet.pth')
    charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)
    
    cfg.merge_from_file('./configs/det/r50_baseline.yaml')
    # manual override some options
    cfg.merge_from_list(["MODEL.DEVICE", "cpu"])

    det_demo = DetDemo(
        cfg,
        min_image_size=800,
        confidence_threshold=0.7,
        output_polygon=True
    )
    
    image = cv2.imread(filepath)
    result_polygons, result_masks, result_boxes = det_demo.run_on_opencv_image(image)
    
    patchs = [image[box[1]:box[3], box[0]:box[2], :] for box in result_boxes]
    patchs = [preprocess(patch, config.dataset_image_width, config.dataset_image_height) for patch in patchs]
    patchs = torch.cat(patchs, dim=0)
    res = model(patchs)
    result_words = postprocess(res, charset, 'alignment')[0]
    
    visual_image = det_demo.visualization(image.copy(), result_polygons, result_masks, result_boxes, result_words)
    
    cv2.imwrite('result.jpg', visual_image)
    return 'result.jpg', pd.DataFrame(result_words)


iface = gr.Interface(
    fn=infer,
    title="张博强毕设展示",
    description="毕设题目：自然场景中任意形状文字的检测与识别\n目前进度：检测",
    inputs=[gr.inputs.Image(label="image", type="filepath")],
    outputs=[gr.outputs.Image(), gr.outputs.Dataframe(headers=['word'])], 
    #examples=['example1.jpg', 'example2.jpg', 'example3.jpg'],
    #article="<a href=\"https://github.com/MhLiao/MaskTextSpotterV3\">GitHub Repo</a>",
).launch(enable_queue=True)