app.py

import os
os.system('pip install --upgrade --no-cache-dir gdown')
os.system('gdown -O ./model_ctw.pth 16qgtD4UOhp0q5e2RYXE1dvuTz_ylZMyb')
#os.system('unzip model_ctw.zip')
os.system('gdown -O ./workdir.zip 10HxLehcJMY9rLd_OyH40HmrySZItuNDt')
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(img):
    filepath = './input.png'
    img.save(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)
    print(image.shape)
    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)
    
    print(visual_image.shape)
    cv2.imwrite('result.jpg', visual_image)
    return ['result.jpg', pd.DataFrame(result_words)]

blocks = gr.Blocks()

input_image = gr.Image(label="image", type="pil")
output_image = gr.Image(label="out_img", type="filepath")
output_word = gr.Dataframe(label="out_word", headers=['word'])

with blocks:
    gr.Markdown('''
    <center><h1 id="title">张博强毕设展示</h1></center>
    <center> 西北工业大学 航海学院本科 张博强 </center>
    <center> 毕设题目：自然场景中任意形状文字的检测与识别 </center>
    <center> 检测：基于<a href="https://github.com/wangyuxin87/ContourNet">ContourNet</a> 识别：基于<a href="https://github.com/FangShancheng/ABINet">ABINet</a> </center>
    ''')
    
    with gr.Row():
        with gr.Column():
            input_image.render()
            button = gr.Button("Submit")
            button.click(fn=infer, inputs=[input_image],
                         outputs=[output_image, output_word],)
        with gr.Column():
            output_image.render()
    with gr.Row():
        output_word.render()
    
    
if __name__ == "__main__":
    blocks.launch(debug=True)
'''
iface = gr.Interface(
    fn=infer,
    title="张博强毕设展示",
    description=description,
    inputs=[gr.inputs.Image(label="image", type="filepath")],
    outputs=[gr.outputs.Image(), gr.outputs.Dataframe(headers=['word'])],
    examples=['figs/test/CANDY.png', 'figs/test/ESPLANADE.png', 'figs/test/KAPPA.png'],
).launch(enable_queue=True, cache_examples=True)
'''