| 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 |
|
|
| |
| |
| 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) |
|
|
| |
| for poly in vertices: |
| draw.line(poly + [poly[0]], fill="red", width=2) |
|
|
| |
| 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) |
|
|
