Update app.py

app.py

@@ -7,39 +7,27 @@ from transformers import pipeline, AutoModelForSequenceClassification, AutoToken
 @st.cache(allow_output_mutation=True)
 def load_models():
     image_pipeline = pipeline("image-to-text", model="microsoft/trocr-large-printed")
-    phishing_model = AutoModelForSequenceClassification.from_pretrained("kithangw/phishing_link_detection", num_labels=2)
+    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")
     return image_pipeline, phishing_model, phishing_tokenizer
 
 image_pipeline, phishing_model, phishing_tokenizer = load_models()
 
-# Define the main function
-def main(image_input):
-    # Convert image to URL text
-    def image2url(image_input):
-        url_for_recognise = image_pipeline(image_input)[0]['generated_text'].replace(" ", "").lower()
-        st.write(f"Recognized URL: {url_for_recognise}")
-        return url_for_recognise
-
-    # Check if the URL text is a phishing link
-    def checkphishing(url_for_recognise):
-        link_token = phishing_tokenizer(url_for_recognise, max_length=512, padding=True, truncation=True, return_tensors='pt')
+# Define the phishing check function
+def check_phishing(url_for_recognize):
+    link_token = phishing_tokenizer(url_for_recognize, max_length=512, padding=True, truncation=True, return_tensors='pt')
 
-        with torch.no_grad():  # Disable gradient calculation for inference
-            output = phishing_model(**link_token)
+    with torch.no_grad():  # Disable gradient calculation for inference
+        output = phishing_model(**link_token)
 
-        probabilities = torch.nn.functional.softmax(output.logits, dim=-1)
-        predicted_class = torch.argmax(probabilities, dim=-1).item()
-        predicted_prob = probabilities[0, predicted_class].item()
+    probabilities = torch.nn.functional.softmax(output.logits, dim=-1)
+    predicted_class = torch.argmax(probabilities, dim=-1).item()
+    predicted_prob = probabilities[0, predicted_class].item()
 
-        labels = ['Not Phishing', 'Phishing']
-        prediction_label = labels[predicted_class]
-        sentence = f"The URL '{url_for_recognise}' is classified as '{prediction_label}' with a probability of {predicted_prob:.2f}."
-        return sentence
-
-    url_text = image2url(image_input)
-    result_sentence = checkphishing(url_text)
-    return result_sentence
+    labels = ['Not Phishing', 'Phishing']
+    prediction_label = labels[predicted_class]
+    sentence = f"The URL '{url_for_recognize}' is classified as '{prediction_label}' with a probability of {predicted_prob:.2f}."
+    return sentence
 
 # Streamlit interface
 st.title("Phishing URL Detection from Image")
@@ -48,6 +36,16 @@ uploaded_image = st.file_uploader("Upload an image of the URL", type=["png", "jp
 if uploaded_image is not None:
     image = Image.open(uploaded_image)
     st.image(image, caption='Uploaded URL Image', use_column_width=True)
-    if st.button('Detect'):
-        result = main(uploaded_image)
-        st.write(result)
+    
+    # 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.")