Spaces:

WofoLab
/

kvc-ocr

Sleeping

Update app.py

2f1623d verified 5 months ago

6.63 kB

"""
KYC OCR — Gradio app for Hugging Face Spaces.
Upload a PAN card, cheque, Aadhaar, or passport (PDF or image) and get
structured field extraction. Runs on CPU using RapidOCR (ONNX).
"""
# ── Patch a known gradio_client bug ──────────────────────────────────────
# gradio_client.utils.get_type assumes schema is always a dict, but JSON
# Schema's `additionalProperties` field can legitimately be a bool. When
# Gradio's /info endpoint runs (HF Spaces health checks ping it), the
# function crashes with: "TypeError: argument of type 'bool' is not iterable".
# Fixed upstream in newer releases but still present in many 5.x. Patch it.
import gradio_client.utils as _gcu
_orig_get_type = _gcu.get_type
def _patched_get_type(schema):
    if not isinstance(schema, dict):
        return "Any"
    return _orig_get_type(schema)
_gcu.get_type = _patched_get_type
# ─────────────────────────────────────────────────────────────────────────

import gradio as gr
import time
import traceback
from pathlib import Path

from kyc_extract import process, RapidOCR

# ── Initialise OCR engine once at startup ─────────────────────────────
print("Loading RapidOCR engine (ONNX models)…")
_t0 = time.time()
ENGINE = RapidOCR()
print(f"Ready. Engine init took {time.time() - _t0:.1f}s")


# ── Pretty-format the extracted fields as Markdown ────────────────────
def _fmt_fields(fields: dict) -> str:
    if not fields:
        return "_No fields extracted._"
    lines = []
    for k, v in fields.items():
        label = k.replace("_", " ").title()
        if isinstance(v, bool):
            v = "✅ Yes" if v else "❌ No"
        elif v is None:
            v = "_(not found)_"
        lines.append(f"**{label}:** {v}")
    return "\n\n".join(lines)


def _fmt_info(result: dict) -> str:
    bits = [
        f"**Detected type:** `{result.get('detected_type', 'unknown').upper()}`",
        f"**OCR time:** {result.get('ocr_time_s')}s",
        f"**Rotation applied:** {result.get('rotation_applied', 0)}°",
        f"**OCR quality score:** {result.get('ocr_score')}",
    ]
    return "  •  ".join(bits)


# ── Main handler ──────────────────────────────────────────────────────
def extract(file):
    if file is None:
        return "_Please upload a file._", "", ""

    file_path = file if isinstance(file, str) else file.name

    try:
        pages = process(file_path, ENGINE)
    except Exception as exc:
        return (
            f"❌ **Error:** {exc}\n\n```\n{traceback.format_exc()}\n```",
            "",
            "",
        )

    if not pages:
        return "_No pages processed._", "", ""

    # Render all pages (cheques and PANs are normally single-page; multi-page
    # PDFs of compound KYC packets get each page rendered separately).
    info_blocks, field_blocks, raw_blocks = [], [], []
    for p in pages:
        if "error" in p:
            info_blocks.append(f"### Page {p['page']}\n_{p['error']}_")
            continue
        info_blocks.append(f"### Page {p['page']}\n{_fmt_info(p)}")
        field_blocks.append(f"### Page {p['page']}\n{_fmt_fields(p['fields'])}")
        raw_blocks.append(
            f"━━━ Page {p['page']} ━━━\n"
            + "\n".join(f"[{c:.2f}] {t}" for c, t in p["_raw_lines"])
        )

    return (
        "\n\n".join(info_blocks),
        "\n\n---\n\n".join(field_blocks) if field_blocks else "_No fields extracted._",
        "\n\n".join(raw_blocks),
    )


# ── Gradio UI ─────────────────────────────────────────────────────────
DESCRIPTION = """
# 🇮🇳 KYC Document OCR
Extract structured fields from Indian KYC documents — **PAN card**, **cheque**,
**Aadhaar**, and **passport** — from images or PDFs. Runs entirely on CPU using
[RapidOCR](https://github.com/RapidAI/RapidOCR) (PaddleOCR PP-OCRv4 models in ONNX).
**Supported documents and fields:**
- **PAN:** PAN number, name, father's name, DOB / date of incorporation, entity type
- **Cheque:** bank name, address, IFSC, account number, MICR code, cheque number, transaction code
- **Aadhaar / Passport:** _coming soon — Aadhaar reads best from QR code (`pyzbar`), passport from MRZ (`passporteye`)_
**Privacy note:** this is a public Space. Don't upload real personal documents.
For production use, fork it and run on private infrastructure.
"""

with gr.Blocks(title="KYC OCR") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="Upload PAN / Cheque (PDF or image)",
                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"],
                type="filepath",
            )
            submit_btn = gr.Button("Extract fields", variant="primary", size="lg")
            # Use an absolute path; skip the Examples block if the example file
            # wasn't uploaded (won't crash the whole app).
            _example_path = Path(__file__).parent / "examples" / "synthetic_pan.png"
            if _example_path.exists():
                gr.Examples(
                    examples=[[str(_example_path)]],
                    inputs=file_input,
                    label="Example (synthetic PAN with fake data)",
                )

        with gr.Column(scale=2):
            info_out = gr.Markdown()
            fields_out = gr.Markdown()

    with gr.Accordion("Raw OCR output (for debugging)", open=False):
        raw_out = gr.Textbox(label="Raw OCR lines with confidence scores",
                             lines=20)

    gr.Markdown(
        "---\n"
        "Pipeline source: `kyc_extract.py`. "
        "All extraction is rule-based — no hardcoded names/banks. "
        "Doc-type detection → optional rotation correction (cheques) → "
        "format-spec regex + heuristics + checksum validation."
    )

    submit_btn.click(
        extract,
        inputs=file_input,
        outputs=[info_out, fields_out, raw_out],
    )
    file_input.change(
        extract,
        inputs=file_input,
        outputs=[info_out, fields_out, raw_out],
    )


if __name__ == "__main__":
    demo.launch()

	"""
	KYC OCR — Gradio app for Hugging Face Spaces.

	Upload a PAN card, cheque, Aadhaar, or passport (PDF or image) and get
	structured field extraction. Runs on CPU using RapidOCR (ONNX).
	"""
	# ── Patch a known gradio_client bug ──────────────────────────────────────
	# gradio_client.utils.get_type assumes schema is always a dict, but JSON
	# Schema's `additionalProperties` field can legitimately be a bool. When
	# Gradio's /info endpoint runs (HF Spaces health checks ping it), the
	# function crashes with: "TypeError: argument of type 'bool' is not iterable".
	# Fixed upstream in newer releases but still present in many 5.x. Patch it.
	import gradio_client.utils as _gcu
	_orig_get_type = _gcu.get_type
	def _patched_get_type(schema):
	if not isinstance(schema, dict):
	return "Any"
	return _orig_get_type(schema)
	_gcu.get_type = _patched_get_type
	# ─────────────────────────────────────────────────────────────────────────

	import gradio as gr
	import time
	import traceback
	from pathlib import Path

	from kyc_extract import process, RapidOCR

	# ── Initialise OCR engine once at startup ─────────────────────────────
	print("Loading RapidOCR engine (ONNX models)…")
	_t0 = time.time()
	ENGINE = RapidOCR()
	print(f"Ready. Engine init took {time.time() - _t0:.1f}s")


	# ── Pretty-format the extracted fields as Markdown ────────────────────
	def _fmt_fields(fields: dict) -> str:
	if not fields:
	return "_No fields extracted._"
	lines = []
	for k, v in fields.items():
	label = k.replace("_", " ").title()
	if isinstance(v, bool):
	v = "✅ Yes" if v else "❌ No"
	elif v is None:
	v = "_(not found)_"
	lines.append(f"{label}: {v}")
	return "\n\n".join(lines)


	def _fmt_info(result: dict) -> str:
	bits = [
	f"Detected type: `{result.get('detected_type', 'unknown').upper()}`",
	f"OCR time: {result.get('ocr_time_s')}s",
	f"Rotation applied: {result.get('rotation_applied', 0)}°",
	f"OCR quality score: {result.get('ocr_score')}",
	]
	return " • ".join(bits)


	# ── Main handler ──────────────────────────────────────────────────────
	def extract(file):
	if file is None:
	return "_Please upload a file._", "", ""

	file_path = file if isinstance(file, str) else file.name

	try:
	pages = process(file_path, ENGINE)
	except Exception as exc:
	return (
	f"❌ Error: {exc}\n\n```\n{traceback.format_exc()}\n```",
	"",
	"",
	)

	if not pages:
	return "_No pages processed._", "", ""

	# Render all pages (cheques and PANs are normally single-page; multi-page
	# PDFs of compound KYC packets get each page rendered separately).
	info_blocks, field_blocks, raw_blocks = [], [], []
	for p in pages:
	if "error" in p:
	info_blocks.append(f"### Page {p['page']}\n_{p['error']}_")
	continue
	info_blocks.append(f"### Page {p['page']}\n{_fmt_info(p)}")
	field_blocks.append(f"### Page {p['page']}\n{_fmt_fields(p['fields'])}")
	raw_blocks.append(
	f"━━━ Page {p['page']} ━━━\n"
	+ "\n".join(f"[{c:.2f}] {t}" for c, t in p["_raw_lines"])
	)

	return (
	"\n\n".join(info_blocks),
	"\n\n---\n\n".join(field_blocks) if field_blocks else "_No fields extracted._",
	"\n\n".join(raw_blocks),
	)


	# ── Gradio UI ─────────────────────────────────────────────────────────
	DESCRIPTION = """
	# 🇮🇳 KYC Document OCR

	Extract structured fields from Indian KYC documents — PAN card, cheque,
	Aadhaar, and passport — from images or PDFs. Runs entirely on CPU using
	[RapidOCR](https://github.com/RapidAI/RapidOCR) (PaddleOCR PP-OCRv4 models in ONNX).

	Supported documents and fields:
	- PAN: PAN number, name, father's name, DOB / date of incorporation, entity type
	- Cheque: bank name, address, IFSC, account number, MICR code, cheque number, transaction code
	- Aadhaar / Passport: _coming soon — Aadhaar reads best from QR code (`pyzbar`), passport from MRZ (`passporteye`)_

	Privacy note: this is a public Space. Don't upload real personal documents.
	For production use, fork it and run on private infrastructure.
	"""

	with gr.Blocks(title="KYC OCR") as demo:
	gr.Markdown(DESCRIPTION)

	with gr.Row():
	with gr.Column(scale=1):
	file_input = gr.File(
	label="Upload PAN / Cheque (PDF or image)",
	file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"],
	type="filepath",
	)
	submit_btn = gr.Button("Extract fields", variant="primary", size="lg")
	# Use an absolute path; skip the Examples block if the example file
	# wasn't uploaded (won't crash the whole app).
	_example_path = Path(__file__).parent / "examples" / "synthetic_pan.png"
	if _example_path.exists():
	gr.Examples(
	examples=[[str(_example_path)]],
	inputs=file_input,
	label="Example (synthetic PAN with fake data)",
	)

	with gr.Column(scale=2):
	info_out = gr.Markdown()
	fields_out = gr.Markdown()

	with gr.Accordion("Raw OCR output (for debugging)", open=False):
	raw_out = gr.Textbox(label="Raw OCR lines with confidence scores",
	lines=20)

	gr.Markdown(
	"---\n"
	"Pipeline source: `kyc_extract.py`. "
	"All extraction is rule-based — no hardcoded names/banks. "
	"Doc-type detection → optional rotation correction (cheques) → "
	"format-spec regex + heuristics + checksum validation."
	)

	submit_btn.click(
	extract,
	inputs=file_input,
	outputs=[info_out, fields_out, raw_out],
	)
	file_input.change(
	extract,
	inputs=file_input,
	outputs=[info_out, fields_out, raw_out],
	)


	if __name__ == "__main__":
	demo.launch()