Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 import json
+import csv
 import os
 import cv2
 import numpy as np
@@ -9,10 +10,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")):
@@ -26,10 +27,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
+# OCR Functions (No changes here)
 def ocr_with_paddle(img):
     ocr = PaddleOCR(lang='en', use_angle_cls=True)
     result = ocr.ocr(img)
@@ -47,22 +48,6 @@ 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:
@@ -79,6 +64,7 @@ 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"
@@ -89,13 +75,14 @@ def generate_ocr(method, img):
     # Perform inference
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = F.softmax(outputs.logits, dim=1)
-        spam_prob = probs[0][1].item()
+        probs = F.softmax(outputs.logits, dim=1)  # Convert logits to probabilities
+        spam_prob = probs[0][1].item()  # Probability of Spam
 
+    # Adjust classification based on threshold (better than argmax)
     label = "Spam" if spam_prob > 0.5 else "Not Spam"
 
-    # Save results to JSON
-    save_to_json(text_output, label)
+    # Save results using external function
+    save_results_to_repo(text_output, label)
 
     return text_output, label
 
@@ -115,5 +102,5 @@ demo = gr.Interface(
 )
 
 # Launch App
-if __name__ == "_main_":
-    demo.launch()
+if __name__ == "__main__":
+    demo.launch()