gradio_app.py

import os
import sys

if "APP_PATH" in os.environ:
    app_path = os.path.abspath(os.environ["APP_PATH"])
    if os.getcwd() != app_path:
        # fix sys.path for import
        os.chdir(app_path)
    if app_path not in sys.path:
        sys.path.append(app_path)

import io
import tempfile
from typing import List

import pypdfium2
import gradio as gr

from surya.models import load_predictors

from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image

from surya.debug.text import draw_text_on_image

from PIL import Image
from surya.recognition.languages import CODE_TO_LANGUAGE, replace_lang_with_code
from surya.table_rec import TableResult
from surya.detection import TextDetectionResult
from surya.recognition import OCRResult
from surya.layout import LayoutResult
from surya.settings import settings
from surya.common.util import rescale_bbox, expand_bbox


# just copy from streamlit_app.py
def run_ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
    from pdftext.extraction import plain_text_output
    with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
        f.write(pdf_file.getvalue())
        f.seek(0)

        # Sample the text from the middle of the PDF
        page_middle = page_count // 2
        page_range = range(max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count))
        text = plain_text_output(f.name, page_range=page_range)

    sample_gap = len(text) // max_samples
    if len(text) == 0 or sample_gap == 0:
        return "This PDF has no text or very little text", ["no text"]

    if sample_gap < sample_len:
        sample_gap = sample_len

    # Split the text into samples for the model
    samples = []
    for i in range(0, len(text), sample_gap):
        samples.append(text[i:i + sample_len])

    results = predictors["ocr_error"](samples)
    label = "This PDF has good text."
    if results.labels.count("bad") / len(results.labels) > .2:
        label = "This PDF may have garbled or bad OCR text."
    return label, results.labels

# just copy from streamlit_app.py
def inline_detection(img) -> (Image.Image, TextDetectionResult):
    text_pred = predictors["detection"]([img])[0]
    text_boxes = [p.bbox for p in text_pred.bboxes]

    inline_pred = predictors["inline_detection"]([img], [text_boxes], include_maps=True)[0]
    inline_polygons = [p.polygon for p in inline_pred.bboxes]
    det_img = draw_polys_on_image(inline_polygons, img.copy(), color='blue')
    return det_img, text_pred, inline_pred

# just copy from streamlit_app.py `name 'inline_pred' is not defined`
def text_detection(img) -> (Image.Image, TextDetectionResult):
    text_pred = predictors["detection"]([img])[0]
    text_polygons = [p.polygon for p in text_pred.bboxes]
    det_img = draw_polys_on_image(text_polygons, img.copy())
    return det_img, text_pred #, inline_pred

# just copy from streamlit_app.py
def layout_detection(img) -> (Image.Image, LayoutResult):
    pred = predictors["layout"]([img])[0]
    polygons = [p.polygon for p in pred.bboxes]
    labels = [f"{p.label}-{p.position}" for p in pred.bboxes]
    layout_img = draw_polys_on_image(polygons, img.copy(), labels=labels, label_font_size=18)
    return layout_img, pred

# just copy from streamlit_app.py
def table_recognition(img, highres_img, skip_table_detection: bool) -> (Image.Image, List[TableResult]):
    if skip_table_detection:
        layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])]
        table_imgs = [highres_img]
    else:
        _, layout_pred = layout_detection(img)
        layout_tables_lowres = [l.bbox for l in layout_pred.bboxes if l.label in ["Table", "TableOfContents"]]
        table_imgs = []
        layout_tables = []
        for tb in layout_tables_lowres:
            highres_bbox = rescale_bbox(tb, img.size, highres_img.size)
            # Slightly expand the box
            highres_bbox = expand_bbox(highres_bbox)
            table_imgs.append(
                highres_img.crop(highres_bbox)
            )
            layout_tables.append(highres_bbox)

    table_preds = predictors["table_rec"](table_imgs)
    table_img = img.copy()

    for results, table_bbox in zip(table_preds, layout_tables):
        adjusted_bboxes = []
        labels = []
        colors = []

        for item in results.cells:
            adjusted_bboxes.append([
                (item.bbox[0] + table_bbox[0]),
                (item.bbox[1] + table_bbox[1]),
                (item.bbox[2] + table_bbox[0]),
                (item.bbox[3] + table_bbox[1])
            ])
            labels.append(item.label)
            if "Row" in item.label:
                colors.append("blue")
            else:
                colors.append("red")
        table_img = draw_bboxes_on_image(adjusted_bboxes, highres_img, labels=labels, label_font_size=18, color=colors)
    return table_img, table_preds

# just copy from streamlit_app.py
def ocr(img, highres_img, langs: List[str]) -> (Image.Image, OCRResult):
    replace_lang_with_code(langs)
    img_pred = predictors["recognition"]([img], [langs], predictors["detection"], highres_images=[highres_img])[0]

    bboxes = [l.bbox for l in img_pred.text_lines]
    text = [l.text for l in img_pred.text_lines]
    rec_img = draw_text_on_image(bboxes, text, img.size, langs)
    return rec_img, img_pred

def open_pdf(pdf_file):
    return pypdfium2.PdfDocument(pdf_file)

def page_counter(pdf_file):
    doc = open_pdf(pdf_file)
    doc_len = len(doc)
    doc.close()
    return doc_len

def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI):
    doc = open_pdf(pdf_file)
    renderer = doc.render(
        pypdfium2.PdfBitmap.to_pil,
        page_indices=[page_num - 1],
        scale=dpi / 72,
    )
    png = list(renderer)[0]
    png_image = png.convert("RGB")
    doc.close()
    return png_image

def get_uploaded_image(in_file):
    return Image.open(in_file).convert("RGB")

# Load models if not already loaded in reload mode
predictors = load_predictors()

with gr.Blocks(title="Surya") as demo:
    gr.Markdown("""
    # Surya OCR Demo

    This app will let you try surya, a multilingual OCR model. It supports text detection + layout analysis in any language, and text recognition in 90+ languages.

    Notes:
    - This works best on documents with printed text.
    - Preprocessing the image (e.g. increasing contrast) can improve results.
    - If OCR doesn't work, try changing the resolution of your image (increase if below 2048px width, otherwise decrease).
    - This supports 90+ languages, see [here](https://github.com/VikParuchuri/surya/tree/master/surya/languages.py) for a full list.

    Find the project [here](https://github.com/VikParuchuri/surya).
    """)

    with gr.Row():
        with gr.Column():
            in_file = gr.File(label="PDF file or image:", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".webp"])
            in_num = gr.Slider(label="Page number", minimum=1, maximum=100, value=1, step=1)
            in_img = gr.Image(label="Select page of Image", type="pil", sources=None)

            text_det_btn = gr.Button("Run Text Detection")
            inline_det_btn = gr.Button("Run Inline Math Detection")
            layout_det_btn = gr.Button("Run Layout Analysis")

            lang_dd = gr.Dropdown(label="Languages", choices=sorted(list(CODE_TO_LANGUAGE.values())), multiselect=True, max_choices=4, info="Select the languages in the image (if known) to improve OCR accuracy.  Optional.")
            text_rec_btn = gr.Button("Run OCR")

            use_pdf_boxes_ckb = gr.Checkbox(label="Use PDF table boxes", value=True, info="Table recognition only: Use the bounding boxes from the PDF file vs text detection model.")
            skip_table_detection_ckb = gr.Checkbox(label="Skip table detection", value=False, info="Table recognition only: Skip table detection and treat the whole image/page as a table.")
            table_rec_btn = gr.Button("Run Table Rec")

            ocr_errors_btn = gr.Button("Run bad PDF text detection")
        with gr.Column():
            result_img = gr.Image(label="Result image")
            result_json = gr.JSON(label="Result json")

        def show_image(file, num=1):
            if file.endswith('.pdf'):
                count = page_counter(file)
                img = get_page_image(file, num, settings.IMAGE_DPI)
                return [
                    gr.update(visible=True, maximum=count),
                    gr.update(value=img)]
            else:
                img = get_uploaded_image(file)
                return [
                    gr.update(visible=False),
                    gr.update(value=img)]

        in_file.upload(
            fn=show_image,
            inputs=[in_file],
            outputs=[in_num, in_img],
        )
        in_num.change(
            fn=show_image,
            inputs=[in_file, in_num],
            outputs=[in_num, in_img],
        )

        # Run Text Detection
        def text_det_img(pil_image):
            det_img, text_pred = text_detection(pil_image)
            return det_img, text_pred.model_dump(exclude=["heatmap", "affinity_map"])
        text_det_btn.click(
            fn=text_det_img,
            inputs=[in_img],
            outputs=[result_img, result_json]
        )
        def inline_det_img(pil_image):
            det_img, text_pred, inline_pred = inline_detection(pil_image)
            json = {
                "text": text_pred.model_dump(exclude=["heatmap", "affinity_map"]),
                "inline": inline_pred.model_dump(exclude=["heatmap", "affinity_map"])
            }
            return det_img, json
        inline_det_btn.click(
            fn=inline_det_img,
            inputs=[in_img],
            outputs=[result_img, result_json]
        )


        # Run layout
        def layout_det_img(pil_image):
            layout_img, pred = layout_detection(pil_image)
            return layout_img, pred.model_dump(exclude=["segmentation_map"])
        layout_det_btn.click(
            fn=layout_det_img,
            inputs=[in_img],
            outputs=[result_img, result_json]
        )
        # Run OCR
        def text_rec_img(pil_image, in_file, page_number, languages):
            if in_file.endswith('.pdf'):
                pil_image_highres = get_page_image(in_file, page_number, dpi=settings.IMAGE_DPI_HIGHRES)
            else:
                pil_image_highres = pil_image
            rec_img, pred = ocr(pil_image, pil_image_highres, languages)
            return rec_img, pred.model_dump()
        text_rec_btn.click(
            fn=text_rec_img,
            inputs=[in_img, in_file, in_num, lang_dd],
            outputs=[result_img, result_json]
        )
        # Run Table Recognition
        def table_rec_img(pil_image, in_file, page_number, use_pdf_boxes, skip_table_detection):
            if in_file.endswith('.pdf'):
                pil_image_highres = get_page_image(in_file, page_number, dpi=settings.IMAGE_DPI_HIGHRES)
            else:
                pil_image_highres = pil_image
            table_img, pred = table_recognition(pil_image, pil_image_highres, skip_table_detection)
            return table_img, [p.model_dump() for p in pred]
        table_rec_btn.click(
            fn=table_rec_img,
            inputs=[in_img, in_file, in_num, use_pdf_boxes_ckb, skip_table_detection_ckb],
            outputs=[result_img, result_json]
        )
        # Run bad PDF text detection
        def ocr_errors_pdf(file, page_count, sample_len=512, max_samples=10, max_pages=15):
            if file.endswith('.pdf'):
                count = page_counter(file)
            else:
                raise gr.Error("This feature only works with PDFs.", duration=5)
            label, results = run_ocr_errors(io.BytesIO(open(file.name, "rb").read()), count)
            return gr.update(label="Result json:" + label, value=results)
        ocr_errors_btn.click(
            fn=ocr_errors_pdf,
            inputs=[in_file, in_num, use_pdf_boxes_ckb, skip_table_detection_ckb],
            outputs=[result_json]
        )

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