Update app.py

app.py

@@ -1,11 +1,11 @@
 import streamlit as st
-import torch
 from PIL import Image
 from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
 
 # Initialize the image-to-text pipeline and models
 @st.cache(allow_output_mutation=True)
 def load_models():
+    # Make sure to use the correct model names and tokenizer
     image_pipeline = pipeline("image-to-text", model="microsoft/trocr-large-printed")
     phishing_model = AutoModelForSequenceClassification.from_pretrained("kithangw/phishing_link_detection", num_labels=2)
     phishing_tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
@@ -31,21 +31,32 @@ def check_phishing(url_for_recognize):
 
 # Streamlit interface
 st.title("Phishing URL Detection from Image")
-uploaded_image = st.file_uploader("Upload an image of the URL", type=["png", "jpg", "jpeg"])
+
+# Text box for URL input
+verified_url = st.text_input("Enter or paste a URL to check for phishing:")
+
+# File uploader to scan the image
+uploaded_image = st.file_uploader("Alternatively, upload an image of the URL", type=["png", "jpg", "jpeg"])
 
 if uploaded_image is not None:
     image = Image.open(uploaded_image)
     st.image(image, caption='Uploaded URL Image', use_column_width=True)
     
-    # Convert image to URL text
-    url_for_recognize = image_pipeline(uploaded_image)[0]['generated_text'].replace(" ", "").lower()
-    st.write("Recognized URL:")
-    # Use a text input to let the user verify and possibly edit the recognized URL
-    verified_url = st.text_input("Verify or edit the recognized URL if necessary:", value=url_for_recognize)
-
-    if st.button('Detect Phishing'):
-        if verified_url:
-            result = check_phishing(verified_url)
-            st.write(result)
-        else:
-            st.write("Please enter a URL to check for phishing.")
+    try:
+        # Process the image with the OCR pipeline
+        ocr_result = image_pipeline(image)[0]['generated_text'].replace(" ", "").lower()
+        # Update the text input with the OCR result
+        st.session_state['verified_url'] = ocr_result
+    except Exception as e:
+        st.error(f"An error occurred during image processing: {e}")
+
+if st.button('Detect Phishing'):
+    if verified_url:
+        result = check_phishing(verified_url)
+        st.write(result)
+    else:
+        st.write("Please enter or upload a URL to check for phishing.")
+
+# Ensure the text box is updated with the OCR result (if any)
+if 'verified_url' in st.session_state and uploaded_image:
+    verified_url = st.session_state['verified_url']