Nischay103 commited on
Commit
b312809
1 Parent(s): da874fb

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +97 -0
app.py ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import torch
4
+ import numpy as np
5
+ from PIL import Image
6
+ import torchvision.transforms as T
7
+ from torch.autograd import Variable
8
+ from huggingface_hub import hf_hub_download
9
+ import streamlit as st
10
+
11
+ # Constants and environment variables
12
+ TOKEN = os.getenv('hf_read_token')
13
+ repo_id = "Nischay103/captcha_recognition"
14
+ model_files = {
15
+ "v1": "captcha_model_v1.pt", "v2": "captcha_model_v2.pt",
16
+ "v3": "captcha_model_v3.pt", "v4": "captcha_model_v4.pt",
17
+ "v5": "captcha_model_v5.pt", "v6": "captcha_model_v6.pt",
18
+ "v7": "captcha_model_v7.pt", "v8": "captcha_model_v8.pt",
19
+ }
20
+
21
+ example_captchas = {
22
+ "v1": "v1/v1_MYCbSs.jpg", "v2": "v2/v2_032891.png", "v3": "v3/v3_5Bg5m.png",
23
+ "v4": "v4/v4_e78d97.jpg", "v5": "v5/v5_DPMTZ.png", "v6": "v6/v6_WD7A.png",
24
+ "v7": "v7/v7_897375.png", "v8": "v8/v8_qxv6x.png",
25
+ }
26
+
27
+ _decode_cls_dims = {
28
+ "63": 63, "37_uppercase": 37,
29
+ "11": 11, "37_lowercase": 37
30
+ }
31
+
32
+ len_dim_pair = {
33
+ "v1": (6, "63"), "v2": (6, "11"), "v3": (5, "63"),
34
+ "v4": (6, "37_lowercase"), "v5": (5, "37_uppercase"),
35
+ "v6": (4, "63"), "v7": (6, "11"), "v8": (5, "37_uppercase")
36
+ }
37
+
38
+ char_sets = {
39
+ "63": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789$",
40
+ "37_uppercase": "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789$",
41
+ "37_lowercase": "abcdefghijklmnopqrstuvwxyz0123456789$",
42
+ "11": "0123456789$"
43
+ }
44
+
45
+ # Load models
46
+ models = {}
47
+ for key, model_file in model_files.items():
48
+ model_path = hf_hub_download(repo_id=repo_id, filename=model_file, token=TOKEN)
49
+ models[key] = torch.jit.load(model_path)
50
+
51
+ # Function to transform image
52
+ def transform_image(image_path):
53
+ transform = T.Compose([T.ToTensor()])
54
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
55
+ image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
56
+ image = Image.fromarray(image)
57
+ image = image.resize((200, 50))
58
+ image = transform(image)
59
+ image = Variable(image).to(device)
60
+ image = image.unsqueeze(1)
61
+ return image
62
+
63
+ # Function to get label from model prediction
64
+ def get_label(model_prediction, model_version):
65
+ max_captcha_len, cls_dim_encoded = len_dim_pair[model_version]
66
+ _cls = char_sets[cls_dim_encoded]
67
+ cls_dim = _decode_cls_dims[cls_dim_encoded]
68
+ lab = ""
69
+ for idx in range(max_captcha_len):
70
+ start = cls_dim * idx
71
+ end = cls_dim * (idx + 1)
72
+ get_char = _cls[torch.argmax(model_prediction[0, start:end])]
73
+ lab += get_char
74
+ return lab
75
+
76
+ # Streamlit app
77
+ st.title("Character Sequence Recognition from Scene-Image (Captcha)")
78
+ st.write("Recognize captchas using different models")
79
+
80
+ uploaded_file = st.file_uploader("Choose a captcha image...", type=["jpg", "png"])
81
+ model_version = st.selectbox("Model Version", list(model_files.keys()), index=0)
82
+
83
+ if uploaded_file is not None:
84
+ # Save uploaded file to a temporary location
85
+ with open("temp_captcha_image.png", "wb") as f:
86
+ f.write(uploaded_file.getbuffer())
87
+
88
+ input_image_path = "temp_captcha_image.png"
89
+ st.image(input_image_path, caption='Uploaded Captcha Image', use_column_width=True)
90
+
91
+ if st.button('Recognize'):
92
+ input = transform_image(input_image_path)
93
+ model = models[model_version]
94
+ with torch.no_grad():
95
+ model_prediction = model(input)
96
+ output = get_label(model_prediction, model_version)
97
+ st.write(f"Recognized Character Sequence: {output}")