Update app.py

app.py

@@ -27,10 +27,10 @@ else:
     model = DistilBertForSequenceClassification.from_pretrained(MODEL_PATH)
     tokenizer = DistilBertTokenizer.from_pretrained(MODEL_PATH)
 
-# Set the model to evaluation mode to disable dropout layers
+# Set model to evaluation mode
 model.eval()
 
-# Load OCR Methods
+# OCR Methods
 def ocr_with_paddle(img):
     ocr = PaddleOCR(lang='en', use_angle_cls=True)
     result = ocr.ocr(img)
@@ -48,8 +48,8 @@ def ocr_with_easy(img):
     results = reader.readtext(gray_image, detail=0)
     return ' '.join(results)
 
-# OCR Function
-def generate_ocr(method, img):
+# OCR Extraction Function
+def extract_text(method, img):
     if img is None:
         raise gr.Error("Please upload an image!")
 
@@ -64,63 +64,58 @@ def generate_ocr(method, img):
     else:  # KerasOCR
         text_output = ocr_with_keras(img)
 
-    # Clean and truncate the extracted text
+    # Clean extracted text
     text_output = text_output.strip()
+    
     if len(text_output) == 0:
-        return "No text detected!", "Cannot classify"
+        return "No text detected!"
 
-    # Classify Text as Spam or Not Spam
+    return text_output
+
+# Classification Function
+def classify_text(text_output):
+    if text_output.strip() == "No text detected!":
+        return text_output, "Cannot classify"
+
+    # Tokenize text
     inputs = tokenizer(text_output, return_tensors="pt", truncation=True, padding=True, max_length=512)
 
+    # Model inference
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = F.softmax(outputs.logits, dim=1)  # Convert logits to probabilities
+        probs = F.softmax(outputs.logits, dim=1)
         prediction = torch.argmax(probs, dim=1).item()
 
     label_map = {0: "Not Spam", 1: "Spam"}
     label = label_map[prediction]
 
-    # Save results using the external save function
+    # Save results automatically
     save_results_to_repo(text_output, label)
 
     return text_output, label
 
-# Save results to JSON file
-RESULTS_JSON = "ocr_results.json"
-
-def save_to_json(text, label):
-    data = {"Extracted Text": text, "Spam Classification": label}
-
-    # Save to JSON file
-    with open(RESULTS_JSON, "w") as json_file:
-        json.dump(data, json_file, indent=4)
-
-    return f"Results saved to {RESULTS_JSON}"
-
-# Create Gradio Interface
+# 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")
-
-# Main OCR Interface
-demo = gr.Interface(
-    fn=generate_ocr,
-    inputs=[method_input, image_input],
-    outputs=[output_text, output_label],
-    title="OCR Spam Classifier",
-    description="Upload an image, extract text, and classify it as Spam or Not Spam.",
-    theme="compact",
-)
-
-# *Attach Save Button to Function*
-save_button.click(
-    fn=save_to_json,
-    inputs=[output_text, output_label],
-    outputs=[save_output]
-)
+method_input = gr.Radio(["PaddleOCR", "EasyOCR", "KerasOCR"], value="PaddleOCR", label="Choose OCR Method")
+output_text = gr.Textbox(label="Extracted Text", interactive=True)
+output_label = gr.Textbox(label="Spam Classification", interactive=False)
+
+# Define UI layout
+with gr.Blocks() as demo:
+    gr.Markdown("## OCR Spam Classifier")
+    
+    with gr.Row():
+        method_input.render()
+    
+    with gr.Row():
+        image_input.render()
+    
+    extract_button = gr.Button("Submit")
+    classify_button = gr.Button("Classify")
+    
+    extract_button.click(fn=extract_text, inputs=[method_input, image_input], outputs=[output_text])
+    classify_button.click(fn=classify_text, inputs=[output_text], outputs=[output_text, output_label])
 
 # Launch App
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()