import torch
import pycocotools

try:
    import detectron2
except:
    import os 
    os.system('pip install git+https://github.com/facebookresearch/detectron2.git@9eb4831')

import os

os.system('pip install git+https://github.com/aim-uofa/AdelaiDet.git')
os.system('pip install Pillow==9.5.0')

import gradio as gr

from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg

from adet.config import defaults

cfg=get_cfg()
cfg.set_new_allowed(True)
cfg.merge_from_file("./configs/attn_R_50.yaml")
cfg.MODEL.WEIGHTS="tt_e2e_attn_R_50.pth"
cfg.MODEL.DEVICE="cpu"

voc_size = cfg.MODEL.BATEXT.VOC_SIZE
CTLABELS = [' ','!','"','#','$','%','&','\'','(',')','*','+',',','-','.','/','0','1','2','3','4','5','6','7','8','9',':',';','<','=','>','?','@','A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z','[','\\',']','^','_','`','a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z','{','|','}','~']

def decode_recognition(rec):
        s = ''
        for c in rec:
            c = int(c)
            if c < voc_size - 1:
                if voc_size == 96:
                    s += CTLABELS[c]
                else:
                    s += str(chr(CTLABELS[c]))
            elif c == voc_size -1:
                s += u'口'
        return s

predictor = DefaultPredictor(cfg)

def inference(image):
    im = image[:,:,::-1]
    outputs = predictor(im)
    recs = outputs["instances"].recs
    scores = outputs["instances"].scores
    i=0
    predictions = []
    for rec in recs:
        transcription = decode_recognition(rec)
        score = f"{scores[i]:.3f}"
        predictions.append(
            {
                "transcription":transcription,
                "score": score
            }
        )
        i+=1
    return predictions

image=gr.Image(type="numpy", label="Input Image")
label=gr.Label()
examples=["./hfs_examples/COLEGIO SD.JPG","./hfs_examples/COMITE GUADALUPANO CHICO.JPG","./hfs_examples/COSAUTLAN VER.JPG"]
intf=gr.Interface(fn=inference, inputs=image, outputs=["text"], examples=examples)
intf.launch(inline=False)