Upload app.py

app.py

@@ -0,0 +1,97 @@
+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
+
+# Constants and environment variables
+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$"
+}
+
+# Load models
+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)
+
+# Function to transform image
+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
+
+# Function to get label from model prediction
+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
+
+# Streamlit app
+st.title("Character Sequence Recognition from Scene-Image (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 Version", list(model_files.keys()), index=0)
+
+if uploaded_file is not None:
+    # Save uploaded file to a temporary location
+    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}")