deveos's picture
Update app.py
bad831b verified
Raw
History Blame
5.95 kB
```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
# -----------------------------
@spaces.GPU
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)
```