Update app.py

app.py

@@ -34,7 +34,7 @@ if not api_key:
     logger.error("GOOGLE_API_KEY not set")
     raise HTTPException(status_code=500, detail="GOOGLE_API_KEY not set")
 genai.configure(api_key=api_key)
-model = genai.GenerativeModel("gemini-2.0-flash")
+model = genai.GenerativeModel("gemini-1.5-pro")
 
 # Set Tesseract path
 pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
@@ -49,28 +49,45 @@ def log_memory_usage():
     mem_info = process.memory_info()
     return f"Memory usage: {mem_info.rss / 1024 / 1024:.2f} MB"
 
-def get_file_hash(pdf_bytes):
-    """Generate MD5 hash of PDF content."""
-    return hashlib.md5(pdf_bytes).hexdigest()
+def get_file_hash(file_bytes):
+    """Generate MD5 hash of file content."""
+    return hashlib.md5(file_bytes).hexdigest()
 
 def get_text_hash(raw_text):
     """Generate MD5 hash of raw text."""
     return hashlib.md5(raw_text.encode('utf-8')).hexdigest()
 
-async def process_page(img, page_idx):
+async def process_image(img_bytes, filename, idx):
+    """Process a single image (JPG/JPEG/PNG) with OCR."""
+    start_time = time.time()
+    logger.info(f"Starting OCR for {filename} image {idx}, {log_memory_usage()}")
+    try:
+        img = Image.open(io.BytesIO(img_bytes))
+        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
+        img_pil = Image.fromarray(cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB))
+        custom_config = r'--oem 1 --psm 6 -l eng+ara'  # Reduced for performance
+        page_text = pytesseract.image_to_string(img_pil, config=custom_config)
+        logger.info(f"Completed OCR for {filename} image {idx}, took {time.time() - start_time:.2f} seconds, {log_memory_usage()}")
+        return page_text + "\n"
+    except Exception as e:
+        logger.error(f"OCR failed for {filename} image {idx}: {str(e)}, {log_memory_usage()}")
+        return ""
+
+async def process_pdf_page(img, page_idx):
     """Process a single PDF page with OCR."""
     start_time = time.time()
-    logger.info(f"Starting OCR for page {page_idx}, {log_memory_usage()}")
+    logger.info(f"Starting OCR for PDF page {page_idx}, {log_memory_usage()}")
     try:
         img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
         gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
         img_pil = Image.fromarray(cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB))
-        custom_config = r'--oem 1 --psm 6 -l eng+ara+hin+spa+ita+rus'  # Reduced for performance
+        custom_config = r'--oem 1 --psm 6 -l eng+ara'  # Reduced for performance
         page_text = pytesseract.image_to_string(img_pil, config=custom_config)
-        logger.info(f"Completed OCR for page {page_idx}, took {time.time() - start_time:.2f} seconds, {log_memory_usage()}")
+        logger.info(f"Completed OCR for PDF page {page_idx}, took {time.time() - start_time:.2f} seconds, {log_memory_usage()}")
         return page_text + "\n"
     except Exception as e:
-        logger.error(f"OCR failed for page {page_idx}: {str(e)}, {log_memory_usage()}")
+        logger.error(f"OCR failed for PDF page {page_idx}: {str(e)}, {log_memory_usage()}")
         return ""
 
 async def process_with_gemini(filename: str, raw_text: str):
@@ -91,17 +108,14 @@ async def process_with_gemini(filename: str, raw_text: str):
 
     try:
         prompt = f"""
-        You are an intelligent invoice data extractor. 
-        Given raw text from an invoice in any language, 
-        extract key business fields in the specified JSON format. 
-        Handle synonyms (e.g., 'total' = 'net', 'tax' = 'GST'/'TDS'/'VAT'). Extract currency from symbol or acronym as well.
-        The 'Products' field is dynamic and may contain multiple items, each with 'qty', 'description', 'unit_price', and 'amount'. 
-        If a field is missing, include it with an empty string ("") or appropriate default (e.g., 0 for numbers).
+        You are an intelligent invoice data extractor. Given raw text from an invoice in any language and extract key business fields in the specified JSON format. Support English. Handle synonyms (e.g., 'total' = 'net', 'tax' = 'GST'/'TDS'). The 'Products' field is dynamic and may contain multiple items, each with 'qty', 'description', 'unit_price', and 'amount'. Detect the currency (e.g., USD, INR, EUR) from symbols ($, ₹, €) or text; default to USD if unclear. If a field is missing, include it with an empty string ("") or appropriate default (e.g., 0 for numbers).
 
-    
+        Raw text:
+        {raw_text}
 
         Output JSON:
         {{
+            "currency": "",
             "Discount_Percentage": "",
             "Due_Date": "",
             "Email_Client": "",
@@ -117,7 +131,6 @@ async def process_with_gemini(filename: str, raw_text: str):
             "invoice date": "",
             "invoice number": "",
             "shipping address": "",
-            "currency": "",
             "total": ""
         }}
         """
@@ -150,33 +163,34 @@ async def extract_and_structure(files: List[UploadFile] = File(...)):
         total_start_time = time.time()
         logger.info(f"Processing file: {file.filename}, {log_memory_usage()}")
 
-        if not file.filename.lower().endswith('.pdf'):
+        # Validate file format
+        valid_extensions = {'.pdf', '.jpg', '.jpeg', '.png'}
+        file_ext = os.path.splitext(file.filename.lower())[1]
+        if file_ext not in valid_extensions:
             fail_count += 1
             output_json["data"].append({
                 "filename": file.filename,
-                "raw_text": "",
-                "structured_data": {"error": "File is not a PDF"},
-                "error": "File is not a PDF"
+                "structured_data": {"error": f"Unsupported file format: {file_ext}"},
+                "error": f"Unsupported file format: {file_ext}"
             })
-            logger.error(f"File {file.filename} is not a PDF")
+            logger.error(f"Unsupported file format for {file.filename}: {file_ext}")
             continue
 
-        # Read PDF into memory
+        # Read file into memory
         try:
-            pdf_start_time = time.time()
-            pdf_bytes = await file.read()
-            pdf_stream = io.BytesIO(pdf_bytes)
-            file_hash = get_file_hash(pdf_bytes)
-            logger.info(f"Read PDF {file.filename}, took {time.time() - pdf_start_time:.2f} seconds, size: {len(pdf_bytes)/1024:.2f} KB, {log_memory_usage()}")
+            file_start_time = time.time()
+            file_bytes = await file.read()
+            file_stream = io.BytesIO(file_bytes)
+            file_hash = get_file_hash(file_bytes)
+            logger.info(f"Read file {file.filename}, took {time.time() - file_start_time:.2f} seconds, size: {len(file_bytes)/1024:.2f} KB, {log_memory_usage()}")
         except Exception as e:
             fail_count += 1
             output_json["data"].append({
                 "filename": file.filename,
-                "raw_text": "",
-                "structured_data": {"error": f"Failed to read PDF: {str(e)}"},
-                "error": f"Failed to read PDF: {str(e)}"
+                "structured_data": {"error": f"Failed to read file: {str(e)}"},
+                "error": f"Failed to read file: {str(e)}"
             })
-            logger.error(f"Failed to read PDF {file.filename}: {str(e)}, {log_memory_usage()}")
+            logger.error(f"Failed to read file {file.filename}: {str(e)}, {log_memory_usage()}")
             continue
 
         # Check raw text cache
@@ -185,44 +199,55 @@ async def extract_and_structure(files: List[UploadFile] = File(...)):
             raw_text = raw_text_cache[file_hash]
             logger.info(f"Raw text cache hit for {file.filename}, {log_memory_usage()}")
         else:
-            # Try extracting embedded text
-            try:
-                extract_start_time = time.time()
-                reader = PdfReader(pdf_stream)
-                for page in reader.pages:
-                    text = page.extract_text()
-                    if text:
-                        raw_text += text + "\n"
-                logger.info(f"Embedded text extraction for {file.filename}, took {time.time() - extract_start_time:.2f} seconds, text length: {len(raw_text)}, {log_memory_usage()}")
-            except Exception as e:
-                logger.warning(f"Embedded text extraction failed for {file.filename}: {str(e)}, {log_memory_usage()}")
-
-            # If no embedded text, perform OCR
-            if not raw_text.strip():
+            if file_ext == '.pdf':
+                # Try extracting embedded text
                 try:
-                    convert_start_time = time.time()
-                    images = convert_from_bytes(pdf_bytes, poppler_path="/usr/local/bin", dpi=100)
-                    logger.info(f"PDF to images conversion for {file.filename}, {len(images)} pages, took {time.time() - convert_start_time:.2f} seconds, {log_memory_usage()}")
+                    extract_start_time = time.time()
+                    reader = PdfReader(file_stream)
+                    for page in reader.pages:
+                        text = page.extract_text()
+                        if text:
+                            raw_text += text + "\n"
+                    logger.info(f"Embedded text extraction for {file.filename}, took {time.time() - extract_start_time:.2f} seconds, text length: {len(raw_text)}, {log_memory_usage()}")
+                except Exception as e:
+                    logger.warning(f"Embedded text extraction failed for {file.filename}: {str(e)}, {log_memory_usage()}")
 
+                # If no embedded text, perform OCR
+                if not raw_text.strip():
+                    try:
+                        convert_start_time = time.time()
+                        images = convert_from_bytes(file_bytes, poppler_path="/usr/local/bin", dpi=100)
+                        logger.info(f"PDF to images conversion for {file.filename}, {len(images)} pages, took {time.time() - convert_start_time:.2f} seconds, {log_memory_usage()}")
+
+                        ocr_start_time = time.time()
+                        page_texts = []
+                        for i, img in enumerate(images):
+                            page_text = await process_pdf_page(img, i)
+                            page_texts.append(page_text)
+                        raw_text = "".join(page_texts)
+                        logger.info(f"Total OCR for {file.filename}, took {time.time() - ocr_start_time:.2f} seconds, text length: {len(raw_text)}, {log_memory_usage()}")
+                    except Exception as e:
+                        fail_count += 1
+                        output_json["data"].append({
+                            "filename": file.filename,
+                            "structured_data": {"error": f"OCR failed: {str(e)}"},
+                            "error": f"OCR failed: {str(e)}"
+                        })
+                        logger.error(f"OCR failed for {file.filename}: {str(e)}, {log_memory_usage()}")
+                        continue
+            else:  # JPG/JPEG/PNG
+                try:
                     ocr_start_time = time.time()
-                    # Sequential processing to reduce memory usage; uncomment for parallel if needed
-                    page_texts = []
-                    for i, img in enumerate(images):
-                        page_text = await process_page(img, i)
-                        page_texts.append(page_text)
-                    # tasks = [process_page(img, i) for i, img in enumerate(images)]
-                    # page_texts = await asyncio.gather(*tasks)
-                    raw_text = "".join(page_texts)
-                    logger.info(f"Total OCR for {file.filename}, took {time.time() - ocr_start_time:.2f} seconds, text length: {len(raw_text)}, {log_memory_usage()}")
+                    raw_text = await process_image(file_bytes, file.filename, 0)
+                    logger.info(f"Image OCR for {file.filename}, took {time.time() - ocr_start_time:.2f} seconds, text length: {len(raw_text)}, {log_memory_usage()}")
                 except Exception as e:
                     fail_count += 1
                     output_json["data"].append({
                         "filename": file.filename,
-                        "raw_text": "",
-                        "structured_data": {"error": f"OCR failed: {str(e)}"},
-                        "error": f"OCR failed: {str(e)}"
+                        "structured_data": {"error": f"Image OCR failed: {str(e)}"},
+                        "error": f"Image OCR failed: {str(e)}"
                     })
-                    logger.error(f"OCR failed for {file.filename}: {str(e)}, {log_memory_usage()}")
+                    logger.error(f"Image OCR failed for {file.filename}: {str(e)}, {log_memory_usage()}")
                     continue
 
             # Normalize text
@@ -240,14 +265,13 @@ async def extract_and_structure(files: List[UploadFile] = File(...)):
         success_count += 1
         output_json["data"].append({
             "filename": file.filename,
-            "raw_text": raw_text,
             "structured_data": structured_data,
             "error": ""
         })
 
         logger.info(f"Total processing for {file.filename}, took {time.time() - total_start_time:.2f} seconds, {log_memory_usage()}")
 
-    output_json["message"] = f"Processed {len(files)} PDFs. {success_count} succeeded, {fail_count} failed."
+    output_json["message"] = f"Processed {len(files)} files. {success_count} succeeded, {fail_count} failed."
     if fail_count > 0 and success_count == 0:
         output_json["success"] = False