Update app.py

app.py

@@ -74,29 +74,27 @@ def ocr_with_tesseract(img):
     return extracted_text, confidences
 
 # OCR & Classification Function
-def generate_ocr(method, image):
-    if image is None:
+def generate_ocr(method, img):
+    if img is None:
         raise gr.Error("Please upload an image!")
 
     # Convert PIL Image to OpenCV format
-    img_cv = preprocess_image(image)
+    img = np.array(img)
 
     # Select OCR method
     if method == "PaddleOCR":
-        extracted_text, confidences = ocr_with_paddle(img_cv)
+        text_output = ocr_with_paddle(img)
     elif method == "EasyOCR":
-        extracted_text, confidences = ocr_with_easy(img_cv)
+        text_output = ocr_with_easy(img)
     elif method == "KerasOCR":
-        extracted_text, confidences = ocr_with_keras(img_cv)
+        text_output = ocr_with_keras(img)
     elif method == "TesseractOCR":
-        extracted_text, confidences = ocr_with_tesseract(img_cv)
+        text_output, _ = ocr_with_tesseract(img)  # Ignore confidence values
     else:
         return "Invalid OCR method", "N/A"
 
-    # Join extracted text into a single string
-    text_output = " ".join(extracted_text).strip()
-
-    # If no text detected, return early
+    # Clean and truncate the extracted text
+    text_output = text_output.strip()
     if len(text_output) == 0:
         return "No text detected!", "Cannot classify"
 
@@ -108,19 +106,15 @@ def generate_ocr(method, image):
         outputs = model(**inputs)
         logits = outputs.logits  # Get raw logits
 
-    # Print raw logits for debugging
+    # Debugging: Print raw logits
     print(f"Raw logits: {logits}")
 
-    # Compare raw logits instead of using softmax
+    # Use raw logits directly instead of softmax
     predicted_class = torch.argmax(logits, dim=1).item()
 
-    print(f"Predicted Class Index: {predicted_class}")  # Debugging output
-
-    # Ensure correct label mapping
-    if predicted_class == 1:
-        label = "Spam"
-    else:
-        label = "Not Spam"
+    # Map class index to labels
+    label_map = {0: "Not Spam", 1: "Spam"}
+    label = label_map.get(predicted_class, "Unknown")
 
     # Save results
     save_results_to_repo(text_output, label)