Update app.py

app.py

@@ -13,7 +13,6 @@ from save_results import save_results_to_repo
 
 # Paths
 MODEL_PATH = "./distilbert_spam_model"
-RESULTS_JSON = "ocr_results.json"
 
 # Ensure model exists
 if not os.path.exists(os.path.join(MODEL_PATH, "pytorch_model.bin")):
@@ -51,7 +50,7 @@ def ocr_with_easy(img):
 # OCR Extraction Function
 def extract_text(method, img):
     if img is None:
-        raise gr.Error("Please upload an image!")
+        return "Error: Please upload an image!", ""
 
     # Convert PIL Image to OpenCV format
     img = np.array(img)
@@ -68,13 +67,13 @@ def extract_text(method, img):
     text_output = text_output.strip()
     
     if len(text_output) == 0:
-        return "No text detected!"
+        return "No text detected!", ""
 
-    return text_output
+    return text_output, ""
 
 # Classification Function
 def classify_text(text_output):
-    if text_output.strip() == "No text detected!":
+    if text_output.strip() in ["No text detected!", "Error: Please upload an image!"]:
         return text_output, "Cannot classify"
 
     # Tokenize text
@@ -95,25 +94,20 @@ def classify_text(text_output):
     return text_output, label
 
 # Gradio Interface
-image_input = gr.Image()
-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()
-    
+
+    method_input = gr.Radio(["PaddleOCR", "EasyOCR", "KerasOCR"], value="PaddleOCR", label="Choose OCR Method")
+    image_input = gr.Image(label="Upload Image")
+
     extract_button = gr.Button("Submit")
     classify_button = gr.Button("Classify")
-    
-    extract_button.click(fn=extract_text, inputs=[method_input, image_input], outputs=[output_text])
+
+    output_text = gr.Textbox(label="Extracted Text", interactive=True)
+    output_label = gr.Textbox(label="Spam Classification", interactive=False)
+
+    # Button Click Bindings
+    extract_button.click(fn=extract_text, inputs=[method_input, image_input], outputs=[output_text, output_label])
     classify_button.click(fn=classify_text, inputs=[output_text], outputs=[output_text, output_label])
 
 # Launch App