Update main.py

main.py

@@ -1,21 +1,16 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 import base64
 from PIL import Image, ImageEnhance
 import pytesseract
 from langdetect import detect, DetectorFactory
-#from googletrans import Translator
 from deep_translator import GoogleTranslator
 import re
 import numpy as np
 import cv2
 import unicodedata
 import io
-import uvicorn
-from fastapi import UploadFile, File
 from pydantic import BaseModel
 
-
-
 # Fix language detection randomness
 DetectorFactory.seed = 0
 
@@ -93,38 +88,20 @@ def extract_field_from_lines(lines, patterns):
 def extract_invoice_fields(text):
     lines = [line.strip() for line in text.split('\n') if line.strip()]
     invoice_number_patterns = [
-      # Tax Invoice with number explicitly mentioned
-      r'(?i)(?:invoice\s*(?:number|no)?\.?\s*[:\-]?\s*)([A-Z0-9][A-Z0-9\-_/]{4,})',
-
-      r'(?i)(?:invoice\s*(?:number|no)?\.?\s*[:\-]?\s*)(?!date)([A-Z0-9][A-Z0-9\-_/]{4,})',
-
-
-
-      # Generic Invoice No. / Invoice #
-      r'(?:invoice\s*(?:number|no|nos|na|#)?\s*[:\-\=\.]?\s*)([A-Z0-9][A-Z0-9\-_/\.]{3,})',
-
-      # Receipt patterns
-      r'(?:receipt\s*(?:number|no|#)?\s*[:\-]?\s*)([A-Z0-9][A-Z0-9\-_/\.]{2,})',
-      
-      # Generic # prefix
-      r'(?:^|\s)#\s*([A-Z0-9][A-Z0-9\-_/\.]{2,})',
-
-      # Order after Receipt
-      r'(?:order\s*)([A-Z0-9][A-Z0-9\-_/\.]{2,})'
-  ]
-
-
+        r'(?i)(?:invoice\s*(?:number|no)?\.?\s*[:\-]?\s*)([A-Z0-9][A-Z0-9\-_/]{4,})',
+        r'(?:invoice\s*(?:number|no|nos|na|#)?\s*[:\-\=\.]?\s*)([A-Z0-9][A-Z0-9\-_/\.]{3,})',
+        r'(?:receipt\s*(?:number|no|#)?\s*[:\-]?\s*)([A-Z0-9][A-Z0-9\-_/\.]{2,})',
+        r'(?:^|\s)#\s*([A-Z0-9][A-Z0-9\-_/\.]{2,})',
+        r'(?:order\s*)([A-Z0-9][A-Z0-9\-_/\.]{2,})'
+    ]
 
-    # Context-aware patterns first (with "date" keywords)
     date_patterns = [
         r'(?:invoice\s*date|bill\s*date|receipt\s*date|date)\s*[:\-]?\s*(\d{1,2}[/-][A-Za-z]{3,9}[/-]?\d{2,4})',
         r'(?:invoice\s*date|bill\s*date|receipt\s*date|date)\s*[:\-]?\s*([A-Za-z]{3,9}[ ]?\d{1,2},?[ ]?\d{4})',
         r'(?:invoice\s*date|bill\s*date|receipt\s*date|date)\s*[:\-]?\s*(\d{4}[/-]\d{1,2}[/-]\d{1,2})',
         r'(?:invoice\s*date|bill\s*date|receipt\s*date|date)\s*[:\-]?\s*(\d{1,2}[/-]\d{1,2}[/-]\d{2,4})',
-        r'(?:receipt\s*date)\s*[:\-]?\s*(\d{1,2}[/-]\d{1,2}[/-]\d{2,4})',
     ]
 
-    # Fallback patterns (no keywords, match only if above fail)
     fallback_date_patterns = [
         r'\b(\d{1,2}\s[A-Za-z]{3,9}\s?\d{2,4})\b',
         r'\b(\d{1,2}[/-][A-Za-z]{3,9}[/-]?\d{2,4})\b',
@@ -133,23 +110,15 @@ def extract_invoice_fields(text):
         r'\b(\d{1,2}[/-]\d{1,2}[/-]\d{2,4})\b',
     ]
 
-
-
-
     amount_patterns = [
         r'(?:total\s*amount|grand\s*total|amount\s*payable|net\s*amount|total|rounding)\s*[:\-]?\s*\₹?\s*([\d,]+\.\d{2})',
-        # r'Total\s+Sales\s*\(Inclusive\s+GST\)\s*[A-Za-z]*\s*([\d,.]+)'
-        #r'[\s:](\d{3,6}\.\d{2})[\s]*$',
-         #r'(?i)(?:total\s*(?:value|due)?|invoice\s*value)\s*[:\-]?\s*(?:₹|Rs\.?|INR)?\s*([\d,.]+)', # Added this pattern
-         r'\b(₹|Rs\.?|INR)\s*([\d,]+\.\d{2})\b', # Added this pattern
-         #r'(?i)(total\s*(amount|value|due)?|invoice\s*value|grand\s*total)[:\-]?\s*(₹|Rs\.?|INR)?\s*([\d,.]+)',
-         r'\b(₹|Rs\.?|INR)\s*([\d,]+\.\d{2})\b'
-
+        r'\b(₹|Rs\.?|INR)\s*([\d,]+\.\d{2})\b',
     ]
 
     invoice_number = extract_field_from_lines(lines, invoice_number_patterns)
     invoice_date = extract_field_from_lines(lines, date_patterns) or extract_field_from_lines(lines, fallback_date_patterns)
     total_amount = extract_field_from_lines(lines, amount_patterns)
+
     if not total_amount:
         numbers = []
         for line in lines:
@@ -157,44 +126,53 @@ def extract_invoice_fields(text):
             numbers += [float(m.replace(',', '')) for m in matches if m]
         if numbers:
             total_amount = f"{max(numbers):.2f}"
+
     return {
         "invoice_number": invoice_number,
         "invoice_date": invoice_date,
         "total_amount": total_amount
     }
 
-# ------------------ API ENDPOINT ------------------
+# ------------------ API ENDPOINTS ------------------
 class ImagePayload(BaseModel):
     image: str
 
+@app.get("/")
+def read_root():
+    return {"status": "ok", "message": "Invoice OCR API is running!"}
+
 @app.post("/predict")
 async def predict(payload: ImagePayload):
-    img_base64 = payload.image
-    if not img_base64:
-        return {"error": "No image provided"}
+    try:
+        img_base64 = payload.image
+        if not img_base64:
+            return {"error": "No image provided"}
 
-    if img_base64.startswith("data:image"):
-        img_base64 = img_base64.split(",")[1]
-    # Preprocess
-    processed_img = preprocess_image(image)
+        # Remove base64 prefix if present
+        if img_base64.startswith("data:image"):
+            img_base64 = img_base64.split(",")[1]
 
-    # OCR + Translation
-    text_data = perform_ocr(processed_img)
+        # Decode base64 to image
+        image_bytes = base64.b64decode(img_base64)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-    # Cleaning
-    cleaned_text = clean_ocr_text(text_data["translated_text"] or text_data["original_text"])
+        # Preprocess
+        processed_img = preprocess_image(image)
 
-    # Extraction
-    fields = extract_invoice_fields(cleaned_text)
+        # OCR + Translation
+        text_data = perform_ocr(processed_img)
 
-    return {
-        "language": text_data["detected_language"],
-        "text": cleaned_text,
-        "fields": fields
-    }
+        # Cleaning
+        cleaned_text = clean_ocr_text(text_data["translated_text"] or text_data["original_text"])
+
+        # Extraction
+        fields = extract_invoice_fields(cleaned_text)
 
+        return {
+            "language": text_data["detected_language"],
+            "text": cleaned_text,
+            "fields": fields
+        }
 
-if __name__ == "__main__":
-    import os
-    port = int(os.environ.get("PORT", 8080))
-    #uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)
+    except Exception as e:
+        return {"error": str(e)}