Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import torch
 import json
-import csv
 import os
 import cv2
 import numpy as np
@@ -15,7 +14,6 @@ from save_results import save_results_to_repo
 # Paths
 MODEL_PATH = "./distilbert_spam_model"
 RESULTS_JSON = "ocr_results.json"
-RESULTS_CSV = "ocr_results.csv"
 
 # Ensure model exists
 if not os.path.exists(os.path.join(MODEL_PATH, "pytorch_model.bin")):
@@ -87,14 +85,23 @@ def generate_ocr(method, img):
 
     return text_output, label
 
+# Save results to JSON file
+def save_to_json(text, label):
+    data = {"extracted_text": text, "classification": label}
+    with open(RESULTS_JSON, "w") as f:
+        json.dump(data, f, indent=4)
+    return "Results saved to JSON file!"
+
 # Gradio Interface
 image_input = gr.Image()
 method_input = gr.Radio(["PaddleOCR", "EasyOCR", "KerasOCR"], value="PaddleOCR")
 output_text = gr.Textbox(label="Extracted Text")
 output_label = gr.Textbox(label="Spam Classification")
+save_button = gr.Button("Save to JSON")
+save_output = gr.Textbox(label="Save Status")
 
 demo = gr.Interface(
-    generate_ocr,
+    fn=generate_ocr,
     inputs=[method_input, image_input],
     outputs=[output_text, output_label],
     title="OCR Spam Classifier",
@@ -102,6 +109,8 @@ demo = gr.Interface(
     theme="compact",
 )
 
-# Launch App
-if __name__ == "__main__":
-    demo.launch()
+# Add Save Button Interaction
+demo.add_component(save_button)
+save_button.click(save_to_json, inputs=[output_text, output_label], outputs=[save_output])
+
+demo.launch()