Spaces:
Running on Zero
Running on Zero
| ```python | |
| import gradio as gr | |
| import json | |
| import re | |
| import spaces | |
| from docling.document_converter import ( | |
| DocumentConverter, | |
| PdfFormatOption | |
| ) | |
| from docling.datamodel.pipeline_options import ( | |
| PdfPipelineOptions, | |
| TesseractCliOcrOptions | |
| ) | |
| from docling.datamodel.base_models import InputFormat | |
| # ----------------------------- | |
| # Helpers | |
| # ----------------------------- | |
| def clean_amount(x): | |
| if not x: | |
| return "" | |
| return x.replace(",", "").strip() | |
| def convert_date(d): | |
| m = re.match(r"(\d{2})-(\d{2})-(\d{2,4})", d) | |
| if not m: | |
| return d | |
| dd, mm, yy = m.groups() | |
| yyyy = "20" + yy if len(yy) == 2 else yy | |
| return f"{yyyy}-{mm}-{dd}" | |
| def detect_voucher_type(text): | |
| text = text.upper() | |
| receipt_words = [ | |
| "CR", | |
| "DEPOSIT", | |
| "SUBSIDY", | |
| "INTEREST", | |
| "CREDIT", | |
| "NEFT", | |
| "IMPS", | |
| "RTGS" | |
| ] | |
| for word in receipt_words: | |
| if word in text: | |
| return "Receipt" | |
| return "Payment" | |
| # ----------------------------- | |
| # Transaction Extraction | |
| # ----------------------------- | |
| def extract_transactions(doc_dict): | |
| transactions = [] | |
| tables = doc_dict.get("tables", []) | |
| for table in tables: | |
| cells = table.get("data", {}).get("table_cells", []) | |
| grouped_rows = {} | |
| for cell in cells: | |
| row_index = cell.get("start_row_offset_idx") | |
| grouped_rows.setdefault(row_index, []).append(cell) | |
| for _, row_cells in sorted(grouped_rows.items()): | |
| row_text = " ".join( | |
| c.get("text", "") for c in row_cells | |
| ) | |
| date_match = re.search( | |
| r"\d{2}-\d{2}-\d{2,4}", | |
| row_text | |
| ) | |
| if not date_match: | |
| continue | |
| amounts = re.findall( | |
| r"\d{1,3}(?:,\d{3})*(?:\.\d{2})|\d+\.\d{2}", | |
| row_text | |
| ) | |
| if len(amounts) < 2: | |
| continue | |
| date = convert_date(date_match.group(0)) | |
| amount = clean_amount(amounts[-2]) | |
| closing_balance = clean_amount(amounts[-1]) | |
| narration = row_text | |
| narration = narration.replace( | |
| date_match.group(0), | |
| "" | |
| ) | |
| for amt in amounts[-2:]: | |
| narration = narration.replace(amt, "") | |
| narration = re.sub( | |
| r"\s+", | |
| " ", | |
| narration | |
| ).strip() | |
| transactions.append({ | |
| "date": date, | |
| "description": narration[:80], | |
| "raw_narration": narration, | |
| "cheque_reference": "", | |
| "voucher_type": detect_voucher_type(narration), | |
| "amount": amount, | |
| "closing_balance": closing_balance | |
| }) | |
| return { | |
| "success": True, | |
| "total_transactions": len(transactions), | |
| "transactions": transactions | |
| } | |
| # ----------------------------- | |
| # Main Convert Function | |
| # ----------------------------- | |
| def convert_document(file, output_format): | |
| if file is None: | |
| return "No file uploaded", {} | |
| try: | |
| pdf_opts = PdfPipelineOptions( | |
| do_ocr=True, | |
| ocr_options=TesseractCliOcrOptions( | |
| lang=["eng"] | |
| ) | |
| ) | |
| converter = DocumentConverter( | |
| format_options={ | |
| InputFormat.PDF: | |
| PdfFormatOption( | |
| pipeline_options=pdf_opts | |
| ) | |
| } | |
| ) | |
| result = converter.convert(file.name) | |
| doc = result.document | |
| # Markdown | |
| if output_format == "Markdown": | |
| converted_text = doc.export_to_markdown() | |
| # Raw JSON | |
| elif output_format == "JSON": | |
| converted_text = json.dumps( | |
| doc.export_to_dict(), | |
| indent=2, | |
| ensure_ascii=False | |
| ) | |
| # Clean Transaction JSON | |
| elif output_format == "Bank Transaction JSON": | |
| tx_json = extract_transactions( | |
| doc.export_to_dict() | |
| ) | |
| converted_text = json.dumps( | |
| tx_json, | |
| indent=2, | |
| ensure_ascii=False | |
| ) | |
| else: | |
| converted_text = "Unsupported format" | |
| metadata = { | |
| "status": "success", | |
| "pages": len( | |
| getattr(doc, "pages", []) | |
| ), | |
| "output_format": output_format | |
| } | |
| return converted_text, metadata | |
| except Exception as e: | |
| return ( | |
| f"ERROR: {str(e)}", | |
| { | |
| "status": "error" | |
| } | |
| ) | |
| # ----------------------------- | |
| # UI | |
| # ----------------------------- | |
| with gr.Blocks() as app: | |
| gr.Markdown( | |
| "# 📄 Docling OCR + Bank Statement Extractor" | |
| ) | |
| gr.Markdown( | |
| "Upload PDF and extract Markdown, JSON, or clean bank transactions." | |
| ) | |
| with gr.Row(): | |
| file_input = gr.File( | |
| label="Upload PDF", | |
| file_types=[".pdf"] | |
| ) | |
| format_input = gr.Radio( | |
| [ | |
| "Markdown", | |
| "JSON", | |
| "Bank Transaction JSON" | |
| ], | |
| value="Bank Transaction JSON", | |
| label="Choose Output Format" | |
| ) | |
| output_text = gr.Textbox( | |
| label="Converted Output", | |
| lines=25 | |
| ) | |
| output_metadata = gr.JSON( | |
| label="Metadata" | |
| ) | |
| convert_button = gr.Button( | |
| "Convert" | |
| ) | |
| convert_button.click( | |
| fn=convert_document, | |
| inputs=[ | |
| file_input, | |
| format_input | |
| ], | |
| outputs=[ | |
| output_text, | |
| output_metadata | |
| ], | |
| api_name="/convert_document" | |
| ) | |
| app.launch(debug=True) | |
| ``` | |