ABCNet / app.py
karensanchez's picture
forcing the use of Pillow 9.5.0
589b4a0
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)