import os
os.system("pip install --upgrade pip")
os.system("pip install torch torchvision")
os.system("pip install numpy==1.21.0 mmcv-full openmim requests")
os.system("mim install mmdet")
os.system('git clone https://github.com/open-mmlab/mmocr.git')
os.system('pip install -r mmocr/requirements.txt')
os.system('pip install -v -e mmocr')
os.system('export PYTHONPATH=$(pwd)/mmocr:$PYTHONPATH')
os.system('pip freeze')
import sys
import requests
sys.path.append('./mmocr')
from gradio import inputs,outputs,Interface
from PIL import Image
from mmocr.utils.ocr import MMOCR

# load image examples
urls = ['https://i0.wp.com/www.kento.com.my/wp-content/uploads/2020/09/ikhsan-sports-shop-3D-stainless-steel-kento.jpg', 'https://automology.s3.ap-southeast-1.amazonaws.com/wp-content/uploads/2019/08/12224246/P.-Mampos.jpg']
for idx, url in enumerate(urls):
  print(url)
  image = Image.open(requests.get(url, stream=True).raw)
  image.save(f"image_{idx}.png")
det = ["DB_r18","DB_r50","DRRG","FCE_IC15","MaskRCNN_CTW","MaskRCNN_IC15","MaskRCNN_IC17","PS_CTW","PS_IC15","TextSnake"]
rec = ["CRNN","SAR","NRTR_1/16-1/8","NRTR_1/8-1/4","RobustScanner","SATRN","SATRN_sm","ABINet","SEG","CRNN_TPS"]
# [MMOCR(config_dir="mmocr/configs/",det=i,recog=None,device="cpu") for i in det]
# [MMOCR(config_dir="mmocr/configs/",det=None,recog=i,device="cpu") for i in rec]

def process_image(detection_model_name,recognizer_model_name,image_path):
    ocr = MMOCR(config_dir="mmocr/configs/",det=detection_model_name,recog=recognizer_model_name,device="cpu")
    results = ocr.readtext(image_path,output='demo/det_out.jpg')
    img = Image.open("demo/det_out.jpg")
    return [img,results[0]["text"]]
title = "Interactive demo: MMOCR"
description = "OpenMMLab Text Detection, Recognition and Understanding Toolbox"
article = "Github repo: https://github.com/open-mmlab/mmocr"
examples =[["TextSnake","RobustScanner",f"image_{i}.png"] for i in range(len(urls))]
#css = """.output_image, .input_image {height: 600px !important}"""
iface = Interface(fn=process_image,
                     inputs=[
                        inputs.Dropdown(det,type="value"),
                        inputs.Dropdown(rec,type="value"),
                        inputs.Image(type="filepath")
                     ],
                     outputs=[outputs.Image(type="pil"),outputs.Textbox()],
                     title=title,
                     description=description,
                     article=article,
                     examples=examples
                    )
iface.launch(debug=False)