Update app.py

app.py

@@ -8,7 +8,6 @@ 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 = {
@@ -42,13 +41,11 @@ char_sets = {
     "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'
@@ -60,7 +57,6 @@ def transform_image(image_path):
     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]
@@ -73,22 +69,20 @@ def get_label(model_prediction, model_version):
         lab += get_char
     return lab
 
-# Streamlit app
-st.title("Character Sequence Recognition from Scene-Image (Captcha)")
-st.write("Recognize captchas using different models")
+st.title("char-seq 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)
+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:
-    # 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)
+    st.image(input_image_path, caption='uploaded captcha image', use_column_width=True)
     
-    if st.button('Recognize'):
+    if st.button('recognize'):
         input = transform_image(input_image_path)
         model = models[model_version]
         with torch.no_grad():
@@ -96,9 +90,9 @@ if uploaded_file is not None:
             output = get_label(model_prediction, model_version)
         st.write(f"Recognized Character Sequence: {output}")
 
-st.write("## Examples")
-cols = st.columns(3)
+st.write("## examples")
+cols = st.columns(4)
 for idx,(model_variant,captcha_path) in enumerate(example_captchas.items()):
-    col = cols[idx % 3]
-    col.image(captcha_path,caption=f'Model {model_variant}: {captcha_path.split("/")[-1]}', width=100)
+    col = cols[idx % 4]
+    col.image(captcha_path,caption=f'{captcha_path.split("/")[-1]}', width=100)
     col.write(f"Model Version: {model_variant}")