main.py

from mmocr.apis import MMOCRInferencer
from mmdet.apis import init_detector, inference_detector
import gradio as gr
import cv2
import sys
import torch

print('Loading model...')
device = 'gpu' if torch.cuda.is_available() else 'cpu'

ocr = MMOCRInferencer(
            det='model/text-det/psenet.py',
            det_weights='model/text-det/psenet.pth',
            rec='model/text-recog/config.py',
            rec_weights='model/text-recog/model.pth',
            device=device)

def get_rec(points):
    xs = []
    ys = []
    for ix, iv in enumerate(points):
        if ix % 2:
            ys.append(iv)
        else:
            xs.append(iv)
    return (min(xs), min(ys)), (max(xs), max(ys))
        
    
def predict(image_input):
    return ocr(image_input)

def run():
    demo = gr.Interface(
        fn=predict,
        inputs=gr.components.Image(),
        outputs=gr.JSON(),
    )

    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    run()