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
@@ -10,10 +9,10 @@ import keras_ocr
 from paddleocr import PaddleOCR
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 import torch.nn.functional as F
-from save_results import save_results_to_repo  # Import the save function
 
 # Paths
 MODEL_PATH = "./distilbert_spam_model"
+RESULTS_JSON = "results.json"
 
 # Ensure model exists
 if not os.path.exists(os.path.join(MODEL_PATH, "pytorch_model.bin")):
@@ -27,10 +26,10 @@ else:
     model = DistilBertForSequenceClassification.from_pretrained(MODEL_PATH)
     tokenizer = DistilBertTokenizer.from_pretrained(MODEL_PATH)
 
-# 🔹 Ensure model is in evaluation mode
+# Ensure model is in evaluation mode
 model.eval()
 
-# OCR Functions (No changes here)
+# OCR Functions
 def ocr_with_paddle(img):
     ocr = PaddleOCR(lang='en', use_angle_cls=True)
     result = ocr.ocr(img)
@@ -48,6 +47,22 @@ def ocr_with_easy(img):
     results = reader.readtext(gray_image, detail=0)
     return ' '.join(results)
 
+# Save results to JSON
+def save_to_json(text, label):
+    data = {"text": text, "classification": label}
+    if os.path.exists(RESULTS_JSON):
+        with open(RESULTS_JSON, "r") as file:
+            try:
+                results = json.load(file)
+            except json.JSONDecodeError:
+                results = []
+    else:
+        results = []
+    
+    results.append(data)
+    with open(RESULTS_JSON, "w") as file:
+        json.dump(results, file, indent=4)
+
 # OCR & Classification Function
 def generate_ocr(method, img):
     if img is None:
@@ -64,25 +79,23 @@ def generate_ocr(method, img):
     else:  # KerasOCR
         text_output = ocr_with_keras(img)
 
-    # 🔹 Preprocess text properly
     text_output = text_output.strip()
     if len(text_output) == 0:
         return "No text detected!", "Cannot classify"
 
-    # 🔹 Tokenize text
+    # Tokenize text
     inputs = tokenizer(text_output, return_tensors="pt", truncation=True, padding=True, max_length=512)
 
-    # 🔹 Perform inference
+    # Perform inference
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = F.softmax(outputs.logits, dim=1)  # Convert logits to probabilities
-        spam_prob = probs[0][1].item()  # Probability of Spam
+        probs = F.softmax(outputs.logits, dim=1)
+        spam_prob = probs[0][1].item()
 
-    # 🔹 Adjust classification based on threshold (better than argmax)
     label = "Spam" if spam_prob > 0.5 else "Not Spam"
 
-    # 🔹 Save results using external function
-    save_results_to_repo(text_output, label)
+    # Save results to JSON
+    save_to_json(text_output, label)
 
     return text_output, label
 
@@ -102,5 +115,5 @@ demo = gr.Interface(
 )
 
 # Launch App
-if __name__ == "__main__":
-    demo.launch()
+if _name_ == "_main_":
+    demo.launch()