add some feature

app.py

@@ -269,7 +269,7 @@ def predict_single_url(url):
     print(f"Processing URL: {url}")
     screenshot_path = take_screenshot(url)
     if not screenshot_path:
-        return f"❌ Error: Unable to capture screenshot for {url}. This may be due to:\n• Too many redirects\n• Website blocking automated access\n• Network connectivity issues\n• Invalid URL", "Screenshot capture failed", None, "", "", "", "", ""
+        return f"❌ Error: Unable to capture screenshot for {url}. This may be due to:\n• Too many redirects\n• Website blocking automated access\n• Network connectivity issues\n• Invalid URL", "Screenshot capture failed", None, "", ""
 
     text = extract_text_from_image(screenshot_path)
     raw_text = text  # Store raw text before cleaning
@@ -290,7 +290,7 @@ def predict_single_url(url):
         confidence = image_probs[0].item() if is_gambling else 1 - image_probs[0].item()
         print(f"[Image-Only] URL: {url}")
         print(f"Prediction: {label} | Confidence: {confidence:.2f}\n")
-        return label, f"Confidence: {confidence:.2f} (Image-Only Model)", screenshot_path, raw_text, "", f"Image Model: {image_probs[0]:.2f}", "", ""
+        return label, f"Confidence: {confidence:.2f} (Image-Only Model)", screenshot_path, raw_text, ""
 
     else:
         clean_text_data = clean_text(text)
@@ -314,15 +314,22 @@ def predict_single_url(url):
         print(f"Text Model Prediction Probability: {text_probs[0]:.2f}")
         print(f"Fusion Final Prediction: {label} | Confidence: {confidence:.2f}\n")
 
-        return label, f"Confidence: {confidence:.2f} (Fusion Model)", screenshot_path, raw_text, clean_text_data, f"Image Model: {image_probs[0]:.2f}", f"Text Model: {text_probs[0]:.2f}", f"Fusion Weights: Image={weights[0]:.2f}, Text={weights[1]:.2f}"
+        return label, f"Confidence: {confidence:.2f} (Fusion Model)", screenshot_path, raw_text, clean_text_data
 
 def predict_batch_urls(file_obj):
     results = []
     content = file_obj.read().decode('utf-8')
     urls = [line.strip() for line in content.splitlines() if line.strip()]
     for url in urls:
-        label, confidence, screenshot_path, raw_text, cleaned_text, image_model_output, text_model_output, fusion_weights = predict_single_url(url)
-        results.append({"url": url, "label": label, "confidence": confidence, "screenshot_path": screenshot_path, "raw_text": raw_text, "cleaned_text": cleaned_text, "image_model_output": image_model_output, "text_model_output": text_model_output, "fusion_weights": fusion_weights})
+        label, confidence, screenshot_path, raw_text, cleaned_text = predict_single_url(url)
+        results.append({
+            "url": url, 
+            "label": label, 
+            "confidence": confidence, 
+            "screenshot_path": screenshot_path, 
+            "raw_text": raw_text, 
+            "cleaned_text": cleaned_text
+        })
 
     df = pd.DataFrame(results)
     print(f"Batch prediction completed for {len(urls)} URLs.")
@@ -338,16 +345,13 @@ with gr.Blocks() as app:
         predict_button = gr.Button("Predict")
         
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=1):
                 label_output = gr.Label()
                 confidence_output = gr.Textbox(label="Confidence", interactive=False)
                 screenshot_output = gr.Image(label="Screenshot", type="filepath")
-            with gr.Column():
-                raw_text_output = gr.Textbox(label="Raw Text", interactive=False)
-                cleaned_text_output = gr.Textbox(label="Cleaned Text", interactive=False)
-                image_model_output = gr.Textbox(label="Image Model Output", interactive=False)
-                text_model_output = gr.Textbox(label="Text Model Output", interactive=False)
-                fusion_weights_output = gr.Textbox(label="Fusion Weights", interactive=False)
+            with gr.Column(scale=1):
+                raw_text_output = gr.Textbox(label="Raw Text", interactive=False, lines=10)
+                cleaned_text_output = gr.Textbox(label="Cleaned Text", interactive=False, lines=10)
 
         predict_button.click(
             fn=predict_single_url, 
@@ -357,10 +361,7 @@ with gr.Blocks() as app:
                 confidence_output, 
                 screenshot_output,
                 raw_text_output,
-                cleaned_text_output,
-                image_model_output,
-                text_model_output,
-                fusion_weights_output
+                cleaned_text_output
             ]
         )