File size: 4,659 Bytes
2c049db 6bfc049 2c049db 7204e43 2c049db 6bfc049 2c049db 6bfc049 2c049db 6bfc049 2c049db 6bfc049 7204e43 2c049db 6bfc049 2c049db 6bfc049 7204e43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | from functools import lru_cache
import numpy as np
import gradio as gr
from PIL import Image, ImageDraw
from paddleocr import PaddleOCR
import bbox_align
# Official PP-OCRv6 tiers. A single selector chooses the matching
# detection + recognition pair.
V6_MODELS = {
"Tiny": ("PP-OCRv6_tiny_det", "PP-OCRv6_tiny_rec"),
"Small": ("PP-OCRv6_small_det", "PP-OCRv6_small_rec"),
"Medium": ("PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
}
@lru_cache(maxsize=len(V6_MODELS))
def get_ocr(tier):
det_model, rec_model = V6_MODELS[tier]
return PaddleOCR(
text_detection_model_name=det_model,
text_recognition_model_name=rec_model,
use_doc_orientation_classify=False,
use_doc_unwarping=False,
use_textline_orientation=False,
device="cpu",
enable_mkldnn=True,
cpu_threads=2,
)
def _result_dict(result_obj):
"""Normalize PaddleOCR's result object to its inner `res` dictionary."""
if isinstance(result_obj, dict):
data = result_obj
elif hasattr(result_obj, "json"):
data = result_obj.json
if callable(data):
data = data()
elif hasattr(result_obj, "res"):
return result_obj.res
else:
raise TypeError(f"Unsupported PaddleOCR result type: {type(result_obj)!r}")
return data.get("res", data)
def _group_lines(polys, width, height):
vertices = [
[(float(x), float(y)) for x, y in np.asarray(poly).reshape(-1, 2)]
for poly in polys
]
page_boundary = [
(0.0, 0.0),
(float(width), 0.0),
(float(width), float(height)),
(0.0, float(height)),
]
groups = bbox_align.process(vertices, page_boundary)
groups = [
sorted(group, key=lambda i: min(p[0] for p in vertices[i]))
for group in groups
]
groups.sort(
key=lambda group: min(min(p[1] for p in vertices[i]) for i in group)
)
return groups, vertices
def _draw_groups(image, groups, vertices):
out = image.copy().convert("RGB")
draw = ImageDraw.Draw(out)
# Thin boxes = OCR detections.
for poly in vertices:
draw.line(poly + [poly[0]], fill="red", width=2)
# Thick boxes = reconstructed logical lines.
for line_no, group in enumerate(groups, start=1):
pts = [p for i in group for p in vertices[i]]
xs = [p[0] for p in pts]
ys = [p[1] for p in pts]
x1, y1, x2, y2 = min(xs), min(ys), max(xs), max(ys)
draw.rectangle((x1, y1, x2, y2), outline="blue", width=4)
draw.text((x1 + 3, max(0, y1 - 14)), str(line_no), fill="blue")
return out
def run_receipt(image, tier):
if image is None:
return None, "Upload a receipt image."
image = image.convert("RGB")
ocr = get_ocr(tier)
results = list(ocr.predict(np.asarray(image)))
if not results:
return image, "No text found."
res = _result_dict(results[0])
texts = list(res.get("rec_texts", []))
scores = list(res.get("rec_scores", []))
polys = list(res.get("rec_polys", []))
items = []
for i, (text, poly) in enumerate(zip(texts, polys)):
text = str(text).strip()
if not text:
continue
score = float(scores[i]) if i < len(scores) else 0.0
items.append((text, score, poly))
if not items:
return image, "No recognized text found."
texts = [x[0] for x in items]
polys = [x[2] for x in items]
groups, vertices = _group_lines(polys, image.width, image.height)
annotated = _draw_groups(image, groups, vertices)
grouped_text = "\n".join(
f"{line_no:02d}: " + " | ".join(texts[i] for i in group)
for line_no, group in enumerate(groups, start=1)
)
return annotated, grouped_text
with gr.Blocks(title="PP-OCRv6 Receipt Line Grouping") as demo:
gr.Markdown(
"# Receipt OCR · PP-OCRv6 + bbox-align\n"
"Upload a receipt, choose a PP-OCRv6 tier, and inspect the reconstructed logical lines."
)
model_tier = gr.Dropdown(
choices=list(V6_MODELS.keys()),
value="Small",
label="PP-OCRv6 model",
)
with gr.Row():
input_image = gr.Image(type="pil", label="Receipt")
output_image = gr.Image(type="pil", label="Grouped lines")
output_text = gr.Textbox(
label="Grouped text",
lines=16,
buttons=["copy"],
)
run_button = gr.Button("Run OCR", variant="primary")
run_button.click(
run_receipt,
inputs=[input_image, model_tier],
outputs=[output_image, output_text],
)
if __name__ == "__main__":
demo.launch(ssr_mode=False)
|