File size: 3,835 Bytes
90505a0 47ddc95 51844ef 90505a0 51844ef 90505a0 51844ef 90505a0 51844ef 90505a0 51844ef 531b32b 2ac0f48 738a12d 5a0ead0 738a12d 531b32b 738a12d 531b32b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import os
import cv2
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as T
from torch.autograd import Variable
from huggingface_hub import hf_hub_download
import streamlit as st
TOKEN = os.getenv('hf_read_token')
repo_id = "Nischay103/captcha_recognition"
model_files = {
"v1": "captcha_model_v1.pt", "v2": "captcha_model_v2.pt",
"v3": "captcha_model_v3.pt", "v4": "captcha_model_v4.pt",
"v5": "captcha_model_v5.pt", "v6": "captcha_model_v6.pt",
"v7": "captcha_model_v7.pt", "v8": "captcha_model_v8.pt",
}
example_captchas = {
"v1": "v1/v1_MYCbSs.jpg", "v2": "v2/v2_032891.png", "v3": "v3/v3_5Bg5m.png",
"v4": "v4/v4_e78d97.jpg", "v5": "v5/v5_DPMTZ.png", "v6": "v6/v6_WD7A.png",
"v7": "v7/v7_897375.png", "v8": "v8/v8_qxv6x.png",
}
_decode_cls_dims = {
"63": 63, "37_uppercase": 37,
"11": 11, "37_lowercase": 37
}
len_dim_pair = {
"v1": (6, "63"), "v2": (6, "11"), "v3": (5, "63"),
"v4": (6, "37_lowercase"), "v5": (5, "37_uppercase"),
"v6": (4, "63"), "v7": (6, "11"), "v8": (5, "37_uppercase")
}
char_sets = {
"63": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789$",
"37_uppercase": "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789$",
"37_lowercase": "abcdefghijklmnopqrstuvwxyz0123456789$",
"11": "0123456789$"
}
models = {}
for key, model_file in model_files.items():
model_path = hf_hub_download(repo_id=repo_id, filename=model_file, token=TOKEN)
models[key] = torch.jit.load(model_path)
def transform_image(image_path):
transform = T.Compose([T.ToTensor()])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
image = Image.fromarray(image)
image = image.resize((200, 50))
image = transform(image)
image = Variable(image).to(device)
image = image.unsqueeze(1)
return image
def get_label(model_prediction, model_version):
max_captcha_len, cls_dim_encoded = len_dim_pair[model_version]
_cls = char_sets[cls_dim_encoded]
cls_dim = _decode_cls_dims[cls_dim_encoded]
lab = ""
for idx in range(max_captcha_len):
start = cls_dim * idx
end = cls_dim * (idx + 1)
get_char = _cls[torch.argmax(model_prediction[0, start:end])]
lab += get_char
return lab
st.title("char-seq recognition from captcha")
st.write("recognize captchas using different models")
uploaded_file = st.file_uploader("choose a captcha image...", type=["jpg", "png"])
model_version = st.selectbox("model variant", list(model_files.keys()), index=0)
if uploaded_file is not None:
with open("temp_captcha_image.png", "wb") as f:
f.write(uploaded_file.getbuffer())
input_image_path = "temp_captcha_image.png"
st.image(input_image_path, caption='uploaded captcha image', use_column_width=True)
if st.button('recognize'):
input = transform_image(input_image_path)
model = models[model_version]
with torch.no_grad():
model_prediction = model(input)
output = get_label(model_prediction, model_version)
st.write(f"Recognized Character Sequence: {output}")
st.write("## examples")
__placeholder__ = st.empty()
cols = st.columns(4)
for idx,(model_variant,captcha_path) in enumerate(example_captchas.items()):
with cols[idx % 4]:
st.image(captcha_path, caption=f'{captcha_path.split("/")[-1]}', width=160)
if st.button(f"Recognize with {model_variant}", key=f"button_{idx}"):
input = transform_image(captcha_path)
model = models[model_variant]
with torch.no_grad():
model_prediction = model(input)
output = get_label(model_prediction, model_variant)
__placeholder__.write(f"Model {model_variant} recognized: {output}") |