|
import os,torch |
|
import gradio as gr |
|
from PIL import Image |
|
from torchvision import transforms as T |
|
from huggingface_hub import hf_hub_download |
|
from tokenizer import Tokenizer |
|
|
|
TOKEN = os.getenv('hf_read_token') |
|
repo_id = "Nischay103/captcha_recognition" |
|
model_file = 'captcha_model.pt' |
|
model_path = hf_hub_download(repo_id=repo_id, filename=model_file, token=TOKEN) |
|
charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" |
|
_tokenizer = Tokenizer(charset) |
|
|
|
def get_transform(img_size=(32,128)): |
|
transforms = [] |
|
transforms.extend([ |
|
T.Resize(img_size, T.InterpolationMode.BICUBIC), |
|
T.ToTensor(), |
|
T.Normalize(0.5, 0.5) |
|
]) |
|
return T.Compose(transforms) |
|
|
|
transform = get_transform() |
|
image = gr.Image(type="pil", label="captcha image") |
|
text = gr.Textbox(label="recognized character sequence") |
|
|
|
def recognize_captcha(input): |
|
input = transform(input.convert('RGB')).unsqueeze(0) |
|
model = torch.jit.load(model_path) |
|
with torch.no_grad(): |
|
probs = model(input) |
|
preds,_ = _tokenizer.decode(probs,beam_width=2) |
|
return preds[0] |
|
|
|
iface = gr.Interface( |
|
fn = recognize_captcha, |
|
inputs = image, |
|
outputs= text, |
|
title = "character sequence recognition from scene-image (captcha)", |
|
description = "recognize captchas ", |
|
examples = ['examples/251615.png','examples/e7dx2r.jpg','examples/980209.png','examples/2DMM4.png', |
|
'examples/H1GQD.png','examples/Sv5Cwm.jpg','examples/Wwx7.png','examples/qw7vr.png', |
|
'examples/mCCyy.jpg','examples/mCGe.jpg','examples/0a0eou.png'] |
|
) |
|
|
|
iface.launch() |